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

@@ -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(), "/")

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
+}