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