Enhance Mihomo WebSocket functionality and testing

- Added support for the Gorilla WebSocket library to improve WebSocket handling in Mihomo.
- Refactored the WebSocket upgrade detection logic to utilize `websocket.IsWebSocketUpgrade`, enhancing reliability.
- Updated the `NewMihomoForward` function to streamline request handling and improve path normalization.
- Introduced new test cases for WebSocket tunnel scenarios, ensuring comprehensive coverage for traffic and memory endpoints.
- Improved error handling and connection management for WebSocket upgrades, ensuring robust communication.

internal/proxy/mihomo_test.go

@@ -40,89 +40,68 @@ func TestMihomoForwardRewritesPathAndAuth(t *testing.T) {
 	assertSameURL(t, cap.got.URL, want)
 }
 
-func TestIsWebSocketUpgrade(t *testing.T) {
-	t.Run("standard headers", func(t *testing.T) {
-		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
+func TestUseMihomoWebSocketTunnel(t *testing.T) {
+	t.Run("gorilla IsWebSocketUpgrade", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
 		r.Header.Set("Connection", "Upgrade")
 		r.Header.Set("Upgrade", "websocket")
-		if !isWebSocketUpgrade(r) {
-			t.Fatal("expected websocket 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")
-		}
-	})
-
-	t.Run("Sec-WebSocket-Key fallback", func(t *testing.T) {
-		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
+		r.Header.Set("Sec-WebSocket-Version", "13")
 		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
-		if !isWebSocketUpgrade(r) {
-			t.Fatal("expected websocket upgrade with Sec-WebSocket-Key")
+		if !useMihomoWebSocketTunnel("traffic", r) {
+			t.Fatal("expected tunnel for full WS handshake")
 		}
 	})
 
-	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("traffic with key only (no Connection token)", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
+		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+		if !useMihomoWebSocketTunnel("traffic", r) {
+			t.Fatal("expected tunnel for traffic+Sec-WebSocket-Key")
 		}
 	})
 
-	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")
+	t.Run("memory with key only", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/memory", nil)
+		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+		if !useMihomoWebSocketTunnel("memory", r) {
+			t.Fatal("expected tunnel for memory+Sec-WebSocket-Key")
+		}
+	})
+
+	t.Run("traffic streaming GET without key", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
+		if useMihomoWebSocketTunnel("traffic", r) {
+			t.Fatal("plain GET /traffic must use HTTP forwarder")
+		}
+	})
+
+	t.Run("proxies path never tunnel", func(t *testing.T) {
+		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/proxies", nil)
+		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
+		if useMihomoWebSocketTunnel("proxies", r) {
+			t.Fatal("proxies must not use WS tunnel")
 		}
 	})
 }
 
-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)
-		}
+func TestMihomoWebSocketURL(t *testing.T) {
+	u, err := url.Parse("http://172.20.0.2:9090/traffic?q=1")
+	if err != nil {
+		t.Fatal(err)
 	}
-	if containsStr(raw, "client-token-must-not-leak") {
-		t.Error("client Authorization leaked into upstream request")
+	s, err := mihomoWebSocketURL(u)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := "ws://172.20.0.2:9090/traffic?q=1"; s != want {
+		t.Fatalf("got %q want %q", s, want)
+	}
+	u2, _ := url.Parse("https://example.com/mem")
+	s2, err := mihomoWebSocketURL(u2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := "wss://example.com/mem"; s2 != want {
+		t.Fatalf("got %q want %q", s2, want)
 	}
 }
-
-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
-}