From 1686840b0e4f938f6aae565ee58b67726facfd35 Mon Sep 17 00:00:00 2001
From: Denozordec <shatskiy.blog@gmail.com>
Date: Tue, 31 Mar 2026 10:22:37 +0700
Subject: [PATCH] Add WebSocket upgrade handling and tests in Mihomo

- Implemented enhanced logic in `isWebSocketUpgrade` to accurately determine WebSocket upgrade requests by checking both "Connection" and "Upgrade" headers.
- Added a new test function `TestIsWebSocketUpgrade` to validate the WebSocket upgrade detection logic, ensuring correct behavior for various header configurations.
---
 internal/proxy/mihomo.go      | 16 +++++++++++++++-
 internal/proxy/mihomo_test.go | 15 +++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/internal/proxy/mihomo.go b/internal/proxy/mihomo.go
index ad9e3c5..d7af141 100644
--- a/internal/proxy/mihomo.go
+++ b/internal/proxy/mihomo.go
@@ -36,7 +36,21 @@ func isWebSocketUpgrade(r *http.Request) bool {
 	if r == nil {
 		return false
 	}
-	return strings.EqualFold(r.Header.Get("Upgrade"), "websocket")
+	// Some clients/proxies pass comma-separated tokens or extra spaces.
+	// Treat request as WS only when both headers contain required upgrade tokens.
+	return headerHasToken(r.Header, "Connection", "upgrade") &&
+		headerHasToken(r.Header, "Upgrade", "websocket")
+}
+
+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(
diff --git a/internal/proxy/mihomo_test.go b/internal/proxy/mihomo_test.go
index a38a7fb..41e1ec1 100644
--- a/internal/proxy/mihomo_test.go
+++ b/internal/proxy/mihomo_test.go
@@ -39,3 +39,18 @@ func TestMihomoForwardRewritesPathAndAuth(t *testing.T) {
 	}
 	assertSameURL(t, cap.got.URL, want)
 }
+
+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")
+	}
+
+	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")
+	}
+}