From 7e88cfcb3e879fb3165da7aa8fdb44a73f41f4aa Mon Sep 17 00:00:00 2001 From: Denozordec Date: Tue, 31 Mar 2026 11:19:40 +0700 Subject: [PATCH] Refactor Mihomo WebSocket handling and enhance tests - Updated the WebSocket upgrade detection logic in `isWebSocketUpgrade` to improve header handling and added a fallback for the "Sec-WebSocket-Key" header. - Refactored the WebSocket proxy logic to use a new `newMihomoWSTunnel` function, enhancing the connection handling process. - Introduced comprehensive test cases in `TestIsWebSocketUpgrade` to validate various WebSocket upgrade scenarios, ensuring robust functionality. - Improved error handling and request building for WebSocket upgrades, ensuring secure and efficient communication. --- internal/proxy/mihomo.go | 128 +++++++++++++++++----------------- internal/proxy/mihomo_test.go | 96 +++++++++++++++++++++---- 2 files changed, 144 insertions(+), 80 deletions(-) diff --git a/internal/proxy/mihomo.go b/internal/proxy/mihomo.go index 617cc52..dd02ba0 100644 --- a/internal/proxy/mihomo.go +++ b/internal/proxy/mihomo.go @@ -16,10 +16,8 @@ import ( // 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 +// REST requests use the same path as NewAliasForward (http.NewRequest + RoundTrip). +// WebSocket (traffic, memory, …) uses raw TCP tunnel: hijack + dial + handshake + bidirectional copy. func NewMihomoForward( target *url.URL, stripPrefix string, @@ -28,7 +26,7 @@ func NewMihomoForward( errHandler func(http.ResponseWriter, *http.Request, error), ) http.Handler { httpH := newAliasForward(target, stripPrefix, "", auth, rt, errHandler, false) - wsH := newMihomoWebSocketReverseProxy(target, stripPrefix, auth, rt, errHandler) + wsH := newMihomoWSTunnel(target, stripPrefix, auth, errHandler) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if isWebSocketUpgrade(r) { wsH.ServeHTTP(w, r) @@ -42,17 +40,14 @@ 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") { + if headerContainsToken(r.Header, "Upgrade", "websocket") && + headerContainsToken(r.Header, "Connection", "upgrade") { return true } return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != "" } -func headerHasToken(h http.Header, key, token string) bool { +func headerContainsToken(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) { @@ -63,11 +58,12 @@ func headerHasToken(h http.Header, key, token string) bool { return false } -func newMihomoWebSocketReverseProxy( +// newMihomoWSTunnel creates a handler that tunnels WebSocket connections to Mihomo +// by hijacking the client TCP connection and building the upstream request from scratch. +func newMihomoWSTunnel( target *url.URL, stripPrefix string, auth string, - rt http.RoundTripper, errHandler func(http.ResponseWriter, *http.Request, error), ) http.Handler { if errHandler == nil { @@ -88,46 +84,36 @@ func newMihomoWebSocketReverseProxy( hj, ok := w.(http.Hijacker) if !ok { - errHandler(w, r, fmt.Errorf("websocket hijack unsupported")) + errHandler(w, r, fmt.Errorf("websocket: hijack not supported by ResponseWriter")) return } - clientConn, clientRW, err := hj.Hijack() + clientConn, clientBuf, err := hj.Hijack() if err != nil { - errHandler(w, r, fmt.Errorf("hijack client conn: %w", err)) + errHandler(w, r, fmt.Errorf("websocket: hijack failed: %w", err)) return } defer clientConn.Close() - upConn, err := dialWebSocketUpstream(dest) + upConn, err := dialUpstream(dest) if err != nil { - _ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway") + _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream connect failed") 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") + reqBytes := buildWSUpgradeRequest(r, dest, auth) + if _, err := upConn.Write(reqBytes); err != nil { + _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream write failed") return } - upBr := bufio.NewReader(upConn) - resp, err := http.ReadResponse(upBr, outReq) + upBuf := bufio.NewReader(upConn) + resp, err := http.ReadResponse(upBuf, nil) if err != nil { - _ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway") + _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream response read failed") return } + if err := resp.Write(clientConn); err != nil { return } @@ -135,18 +121,48 @@ func newMihomoWebSocketReverseProxy( 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 + done := make(chan struct{}, 2) + go func() { _, _ = io.Copy(upConn, clientBuf); done <- struct{}{} }() + go func() { _, _ = io.Copy(clientConn, upBuf); done <- struct{}{} }() + <-done }) } -func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) { +// buildWSUpgradeRequest constructs a raw HTTP/1.1 WebSocket upgrade request +// with only the headers required by RFC 6455 + Authorization for Mihomo. +// This avoids any extra headers that Go's Request.Write may add. +func buildWSUpgradeRequest(orig *http.Request, dest *url.URL, auth string) []byte { + reqURI := dest.RequestURI() + if reqURI == "" { + reqURI = "/" + } + var b strings.Builder + fmt.Fprintf(&b, "GET %s HTTP/1.1\r\n", reqURI) + fmt.Fprintf(&b, "Host: %s\r\n", dest.Host) + b.WriteString("Connection: Upgrade\r\n") + b.WriteString("Upgrade: websocket\r\n") + if v := orig.Header.Get("Sec-WebSocket-Version"); v != "" { + fmt.Fprintf(&b, "Sec-WebSocket-Version: %s\r\n", v) + } + if v := orig.Header.Get("Sec-WebSocket-Key"); v != "" { + fmt.Fprintf(&b, "Sec-WebSocket-Key: %s\r\n", v) + } + if v := orig.Header.Get("Sec-WebSocket-Protocol"); v != "" { + fmt.Fprintf(&b, "Sec-WebSocket-Protocol: %s\r\n", v) + } + if v := orig.Header.Get("Sec-WebSocket-Extensions"); v != "" { + fmt.Fprintf(&b, "Sec-WebSocket-Extensions: %s\r\n", v) + } + if auth != "" { + fmt.Fprintf(&b, "Authorization: %s\r\n", auth) + } + b.WriteString("\r\n") + return []byte(b.String()) +} + +func dialUpstream(dest *url.URL) (net.Conn, error) { if dest == nil { - return nil, fmt.Errorf("nil destination") + return nil, errors.New("nil destination URL") } hostPort := dest.Host if !strings.Contains(hostPort, ":") { @@ -166,7 +182,7 @@ func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) { } } -func writeRawHTTPError(conn net.Conn, code int, text string) error { +func rawHTTPError(conn net.Conn, code int, text string) error { reason := http.StatusText(code) if reason == "" { reason = "Error" @@ -174,30 +190,12 @@ func writeRawHTTPError(conn net.Conn, code int, text string) 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) + _, 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(), "/") diff --git a/internal/proxy/mihomo_test.go b/internal/proxy/mihomo_test.go index 65ce798..b76f284 100644 --- a/internal/proxy/mihomo_test.go +++ b/internal/proxy/mihomo_test.go @@ -41,22 +41,88 @@ func TestMihomoForwardRewritesPathAndAuth(t *testing.T) { } func TestIsWebSocketUpgrade(t *testing.T) { - r1 := httptest.NewRequest(http.MethodGet, "http://gw/api/mtg/mihomo/traffic", nil) - r1.Header.Set("Connection", "keep-alive, Upgrade") - r1.Header.Set("Upgrade", "websocket") - if !isWebSocketUpgrade(r1) { - t.Fatal("expected websocket upgrade for tokenized headers") - } + t.Run("standard headers", func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/ws", nil) + r.Header.Set("Connection", "Upgrade") + r.Header.Set("Upgrade", "websocket") + if !isWebSocketUpgrade(r) { + t.Fatal("expected websocket upgrade") + } + }) - r2 := httptest.NewRequest(http.MethodGet, "http://gw/api/mtg/mihomo/traffic", nil) - r2.Header.Set("Upgrade", "websocket") - if isWebSocketUpgrade(r2) { - t.Fatal("expected non-websocket when Connection lacks upgrade") - } + t.Run("tokenized Connection", func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/ws", nil) + r.Header.Set("Connection", "keep-alive, Upgrade") + r.Header.Set("Upgrade", "websocket") + if !isWebSocketUpgrade(r) { + t.Fatal("expected websocket upgrade for tokenized Connection") + } + }) - r3 := httptest.NewRequest(http.MethodGet, "http://gw/api/mtg/mihomo/traffic", nil) - r3.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") - if !isWebSocketUpgrade(r3) { - t.Fatal("expected websocket upgrade when Sec-WebSocket-Key is present") + t.Run("Sec-WebSocket-Key fallback", func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/ws", nil) + r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + if !isWebSocketUpgrade(r) { + t.Fatal("expected websocket upgrade with Sec-WebSocket-Key") + } + }) + + t.Run("no upgrade headers", func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/ws", nil) + if isWebSocketUpgrade(r) { + t.Fatal("plain GET should not be detected as websocket") + } + }) + + t.Run("Upgrade without Connection", func(t *testing.T) { + r := httptest.NewRequest(http.MethodGet, "/ws", nil) + r.Header.Set("Upgrade", "websocket") + if isWebSocketUpgrade(r) { + t.Fatal("should not match without Connection header or Sec-WebSocket-Key") + } + }) +} + +func TestBuildWSUpgradeRequest(t *testing.T) { + orig := httptest.NewRequest(http.MethodGet, "http://gw/api/mtg/mihomo/traffic", nil) + orig.Header.Set("Connection", "Upgrade") + orig.Header.Set("Upgrade", "websocket") + orig.Header.Set("Sec-WebSocket-Version", "13") + orig.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + orig.Header.Set("Authorization", "Bearer client-token-must-not-leak") + + dest, _ := url.Parse("http://172.20.0.2:9090/traffic") + raw := string(buildWSUpgradeRequest(orig, dest, "Bearer upstream-secret")) + + for _, want := range []string{ + "GET /traffic HTTP/1.1\r\n", + "Host: 172.20.0.2:9090\r\n", + "Connection: Upgrade\r\n", + "Upgrade: websocket\r\n", + "Sec-WebSocket-Version: 13\r\n", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n", + "Authorization: Bearer upstream-secret\r\n", + "\r\n", + } { + if !containsStr(raw, want) { + t.Errorf("request missing %q\ngot:\n%s", want, raw) + } + } + if containsStr(raw, "client-token-must-not-leak") { + t.Error("client Authorization leaked into upstream request") } } + +func containsStr(s, sub string) bool { + return len(s) >= len(sub) && (s == sub || len(sub) == 0 || + (len(s) > 0 && len(sub) > 0 && stringContains(s, sub))) +} + +func stringContains(s, sub string) bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +}