telemt-api/internal/proxy/mihomo.go

package proxy

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// 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
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		NormalizeRequestURLPath(r)
		p := r.URL.Path
		if !strings.HasPrefix(p, stripPrefix) {
			errHandler(w, r, fmt.Errorf("path %q: missing strip prefix %q", p, stripPrefix))
			return
		}
		rest := strings.TrimPrefix(strings.TrimPrefix(p, stripPrefix), "/")
		dest := JoinPathPrefix(target, "/", rest)
		du := *dest
		du.RawQuery = r.URL.RawQuery
		dest = &du

		hj, ok := w.(http.Hijacker)
		if !ok {
			errHandler(w, r, fmt.Errorf("websocket hijack unsupported"))
			return
		}
		clientConn, clientRW, err := hj.Hijack()
		if err != nil {
			errHandler(w, r, fmt.Errorf("hijack client conn: %w", err))
			return
		}
		defer clientConn.Close()

		upConn, err := dialWebSocketUpstream(dest)
		if err != nil {
			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
			return
		}
		defer upConn.Close()

		outReq := r.Clone(r.Context())
		outReq.URL = dest
		outReq.Host = dest.Host
		outReq.RequestURI = ""
		outReq.Proto = "HTTP/1.1"
		outReq.ProtoMajor = 1
		outReq.ProtoMinor = 1
		outReq.Header = cloneHeader(r.Header)
		outReq.Header.Del("Authorization")
		if auth != "" {
			outReq.Header.Set("Authorization", auth)
		}
		if err := outReq.Write(upConn); err != nil {
			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
			return
		}

		upBr := bufio.NewReader(upConn)
		resp, err := http.ReadResponse(upBr, outReq)
		if err != nil {
			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
			return
		}
		if err := resp.Write(clientConn); err != nil {
			return
		}
		if resp.StatusCode != http.StatusSwitchingProtocols {
			return
		}

		clientSrc := io.MultiReader(clientRW, clientConn)
		upSrc := io.MultiReader(upBr, upConn)
		errCh := make(chan error, 2)
		go proxyCopy(errCh, upConn, clientSrc)
		go proxyCopy(errCh, clientConn, upSrc)
		<-errCh
	})
}

func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) {
	if dest == nil {
		return nil, fmt.Errorf("nil destination")
	}
	hostPort := dest.Host
	if !strings.Contains(hostPort, ":") {
		switch strings.ToLower(dest.Scheme) {
		case "https", "wss":
			hostPort += ":443"
		default:
			hostPort += ":80"
		}
	}
	d := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
	switch strings.ToLower(dest.Scheme) {
	case "https", "wss":
		return tls.DialWithDialer(d, "tcp", hostPort, &tls.Config{ServerName: dest.Hostname()})
	default:
		return d.Dial("tcp", hostPort)
	}
}

func writeRawHTTPError(conn net.Conn, code int, text string) error {
	reason := http.StatusText(code)
	if reason == "" {
		reason = "Error"
	}
	if text == "" {
		text = reason
	}
	_, err := fmt.Fprintf(conn, "HTTP/1.1 %d %s\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", code, reason, len(text), text)
	return err
}

func proxyCopy(errCh chan<- error, dst io.Writer, src io.Reader) {
	_, err := io.Copy(dst, src)
	if err != nil && !isNetClosed(err) {
		errCh <- err
		return
	}
	errCh <- nil
}

func isNetClosed(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, net.ErrClosed) {
		return true
	}
	msg := err.Error()
	return strings.Contains(msg, "use of closed network connection")
}

// 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},
	})
}