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.
This commit is contained in:
@@ -36,7 +36,21 @@ func isWebSocketUpgrade(r *http.Request) bool {
|
|||||||
if r == nil {
|
if r == nil {
|
||||||
return false
|
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(
|
func newMihomoWebSocketReverseProxy(
|
||||||
|
|||||||
@@ -39,3 +39,18 @@ func TestMihomoForwardRewritesPathAndAuth(t *testing.T) {
|
|||||||
}
|
}
|
||||||
assertSameURL(t, cap.got.URL, want)
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user