telemt-api/internal/proxy/mihomo.go

package proxy

import (
	"encoding/json"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

// NewMihomoForward proxies /api/{alias}/mihomo/... to Mihomo external-controller.
//
// REST и обычные GET/POST идут тем же путём, что и NewAliasForward (http.NewRequest + RoundTrip):
// httputil.ReverseProxy для всего Mihomo часто давал 400 на upstream при том, что curl к контроллеру работал.
// Только WebSocket (traffic, memory, …) остаётся на ReverseProxy + Rewrite.
// Чеклист при повторении проблемы: docs/GATEWAY_RUN.md#mihomo-debug-400
func NewMihomoForward(
	target *url.URL,
	stripPrefix string,
	auth string,
	rt http.RoundTripper,
	errHandler func(http.ResponseWriter, *http.Request, error),
) http.Handler {
	httpH := newAliasForward(target, stripPrefix, "", auth, rt, errHandler, false)
	wsH := newMihomoWebSocketReverseProxy(target, stripPrefix, auth, rt, errHandler)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if isWebSocketUpgrade(r) {
			wsH.ServeHTTP(w, r)
			return
		}
		httpH.ServeHTTP(w, r)
	})
}

func isWebSocketUpgrade(r *http.Request) bool {
	if r == nil {
		return false
	}
	// Be tolerant to proxy/header quirks:
	// - RFC path: Connection: upgrade + Upgrade: websocket
	// - Fallback: Sec-WebSocket-Key presence strongly indicates WS handshake.
	if headerHasToken(r.Header, "Upgrade", "websocket") &&
		headerHasToken(r.Header, "Connection", "upgrade") {
		return true
	}
	return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != ""
}

func headerHasToken(h http.Header, key, token string) bool {
	for _, v := range h.Values(key) {
		for _, part := range strings.Split(v, ",") {
			if strings.EqualFold(strings.TrimSpace(part), token) {
				return true
			}
		}
	}
	return false
}

func newMihomoWebSocketReverseProxy(
	target *url.URL,
	stripPrefix string,
	auth string,
	rt http.RoundTripper,
	errHandler func(http.ResponseWriter, *http.Request, error),
) http.Handler {
	if errHandler == nil {
		errHandler = defaultForwardErrorHandler
	}
	if rt == nil {
		rt = http.DefaultTransport
	}
	rp := &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			NormalizeRequestURLPath(pr.Out)
			p := pr.Out.URL.Path
			if !strings.HasPrefix(p, stripPrefix) {
				return
			}
			rest := strings.TrimPrefix(strings.TrimPrefix(p, stripPrefix), "/")
			dest := JoinPathPrefix(target, "/", rest)
			du := *dest
			du.RawQuery = pr.Out.URL.RawQuery
			out := pr.Out
			out.URL = &du
			out.Header.Del("Host")
			out.Host = du.Host
			out.RequestURI = ""
			out.Proto = "HTTP/1.1"
			out.ProtoMajor = 1
			out.ProtoMinor = 1
			out.Header.Del("Authorization")
			if auth != "" {
				out.Header.Set("Authorization", auth)
			}
		},
		Transport: rt,
		FlushInterval: -1,
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			errHandler(w, r, err)
		},
	}
	return rp
}

// MihomoMetaJSON returns a JSON body for GET .../mihomo/meta (display URL without credentials).
func MihomoMetaJSON(target *url.URL) []byte {
	display := strings.TrimSuffix(target.String(), "/")
	b, _ := json.Marshal(map[string]any{
		"ok":              true,
		"controller_base": display,
	})
	return b
}

// MihomoJSONError writes a JSON error for Mihomo routes when proxy is not configured.
func MihomoJSONError(w http.ResponseWriter, code, msg string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusNotFound)
	_ = json.NewEncoder(w).Encode(map[string]any{
		"ok":    false,
		"error": map[string]string{"code": code, "message": msg},
	})
}