From 49f36d2b8b4c27f343b71c62150cead9caec1c79 Mon Sep 17 00:00:00 2001 From: Denozordec Date: Tue, 31 Mar 2026 11:41:56 +0700 Subject: [PATCH] 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. --- go.mod | 1 + internal/proxy/mihomo.go | 207 +++++++++++++++------------------- internal/proxy/mihomo_test.go | 119 ++++++++----------- 3 files changed, 142 insertions(+), 185 deletions(-) diff --git a/go.mod b/go.mod index 727332b..524e30a 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/telemt/telemt-api go 1.22 require ( + github.com/gorilla/websocket v1.5.3 github.com/oschwald/geoip2-golang v1.11.0 github.com/prometheus/client_golang v1.20.5 gopkg.in/yaml.v3 v3.0.1 diff --git a/internal/proxy/mihomo.go b/internal/proxy/mihomo.go index dd02ba0..fb370c0 100644 --- a/internal/proxy/mihomo.go +++ b/internal/proxy/mihomo.go @@ -1,23 +1,21 @@ package proxy import ( - "bufio" "crypto/tls" "encoding/json" - "errors" "fmt" - "io" - "net" "net/http" "net/url" "strings" "time" + + "github.com/gorilla/websocket" ) // NewMihomoForward proxies /api/{alias}/mihomo/... to Mihomo external-controller. // -// REST requests use the same path as NewAliasForward (http.NewRequest + RoundTrip). -// WebSocket (traffic, memory, …) uses raw TCP tunnel: hijack + dial + handshake + bidirectional copy. +// REST: same path as NewAliasForward (http.NewRequest + RoundTrip). +// WebSocket (traffic, memory, …): gorilla/websocket Dial to upstream, then Upgrader on client, then frame relay. func NewMihomoForward( target *url.URL, stripPrefix string, @@ -28,7 +26,14 @@ func NewMihomoForward( httpH := newAliasForward(target, stripPrefix, "", auth, rt, errHandler, false) wsH := newMihomoWSTunnel(target, stripPrefix, auth, errHandler) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if isWebSocketUpgrade(r) { + NormalizeRequestURLPath(r) + p := r.URL.Path + if !strings.HasPrefix(p, stripPrefix) { + httpH.ServeHTTP(w, r) + return + } + rest := strings.TrimPrefix(strings.TrimPrefix(p, stripPrefix), "/") + if useMihomoWebSocketTunnel(rest, r) { wsH.ServeHTTP(w, r) return } @@ -36,30 +41,54 @@ func NewMihomoForward( }) } -func isWebSocketUpgrade(r *http.Request) bool { +// useMihomoWebSocketTunnel is true when the client is doing a WS handshake. +// For /traffic and /memory we also accept Sec-WebSocket-Key alone: some hops strip +// Connection/Upgrade but leave Sec-WebSocket-Key; plain streaming GET has no key. +func useMihomoWebSocketTunnel(rest string, r *http.Request) bool { if r == nil { return false } - if headerContainsToken(r.Header, "Upgrade", "websocket") && - headerContainsToken(r.Header, "Connection", "upgrade") { + if websocket.IsWebSocketUpgrade(r) { return true } - return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != "" -} - -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) { - return true - } - } + if r.Method != http.MethodGet { + return false + } + switch rest { + case "traffic", "memory": + return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != "" + default: + return false } - return false } -// newMihomoWSTunnel creates a handler that tunnels WebSocket connections to Mihomo -// by hijacking the client TCP connection and building the upstream request from scratch. +func mihomoWebSocketURL(dest *url.URL) (string, error) { + if dest == nil { + return "", fmt.Errorf("nil destination URL") + } + u := *dest + switch strings.ToLower(u.Scheme) { + case "http": + u.Scheme = "ws" + case "https": + u.Scheme = "wss" + default: + return "", fmt.Errorf("unsupported mihomo scheme %q", dest.Scheme) + } + return u.String(), nil +} + +func tlsConfigForMihomoWS(dest *url.URL) *tls.Config { + if dest == nil || strings.ToLower(dest.Scheme) != "https" { + return nil + } + return &tls.Config{ + ServerName: dest.Hostname(), + MinVersion: tls.VersionTLS12, + } +} + +// newMihomoWSTunnel: dial upstream WebSocket first (with Authorization), then upgrade the client, then relay frames. func newMihomoWSTunnel( target *url.URL, stripPrefix string, @@ -82,118 +111,66 @@ func newMihomoWSTunnel( du.RawQuery = r.URL.RawQuery dest = &du - hj, ok := w.(http.Hijacker) - if !ok { - errHandler(w, r, fmt.Errorf("websocket: hijack not supported by ResponseWriter")) - return - } - clientConn, clientBuf, err := hj.Hijack() + wsURL, err := mihomoWebSocketURL(dest) if err != nil { - errHandler(w, r, fmt.Errorf("websocket: hijack failed: %w", err)) + errHandler(w, r, err) return } - defer clientConn.Close() - upConn, err := dialUpstream(dest) + dialHdr := make(http.Header) + if auth != "" { + dialHdr.Set("Authorization", auth) + } + d := websocket.Dialer{ + HandshakeTimeout: 15 * time.Second, + Proxy: func(*http.Request) (*url.URL, error) { return nil, nil }, + TLSClientConfig: tlsConfigForMihomoWS(dest), + } + upConn, _, err := d.Dial(wsURL, dialHdr) if err != nil { - _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream connect failed") + errHandler(w, r, fmt.Errorf("mihomo ws dial: %w", err)) return } defer upConn.Close() - reqBytes := buildWSUpgradeRequest(r, dest, auth) - if _, err := upConn.Write(reqBytes); err != nil { - _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream write failed") - return + // Upgrader requires Connection+Upgrade; some proxies strip them but leave Sec-WebSocket-Key. + if !websocket.IsWebSocketUpgrade(r) { + r.Header.Set("Connection", "Upgrade") + r.Header.Set("Upgrade", "websocket") + if strings.TrimSpace(r.Header.Get("Sec-WebSocket-Version")) == "" { + r.Header.Set("Sec-WebSocket-Version", "13") + } } - upBuf := bufio.NewReader(upConn) - resp, err := http.ReadResponse(upBuf, nil) + upgrader := websocket.Upgrader{ + HandshakeTimeout: 15 * time.Second, + CheckOrigin: func(*http.Request) bool { return true }, + } + clientConn, err := upgrader.Upgrade(w, r, nil) if err != nil { - _ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream response read failed") return } + defer clientConn.Close() - if err := resp.Write(clientConn); err != nil { - return - } - if resp.StatusCode != http.StatusSwitchingProtocols { - return - } - - done := make(chan struct{}, 2) - go func() { _, _ = io.Copy(upConn, clientBuf); done <- struct{}{} }() - go func() { _, _ = io.Copy(clientConn, upBuf); done <- struct{}{} }() - <-done + errCh := make(chan error, 2) + go relayWSFrames(upConn, clientConn, errCh) + go relayWSFrames(clientConn, upConn, errCh) + <-errCh }) } -// 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, errors.New("nil destination URL") - } - hostPort := dest.Host - if !strings.Contains(hostPort, ":") { - switch strings.ToLower(dest.Scheme) { - case "https", "wss": - hostPort += ":443" - default: - hostPort += ":80" +func relayWSFrames(dst, src *websocket.Conn, errCh chan<- error) { + for { + mt, data, err := src.ReadMessage() + if err != nil { + errCh <- err + return + } + if err := dst.WriteMessage(mt, data); err != nil { + errCh <- err + return } } - d := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} - switch strings.ToLower(dest.Scheme) { - case "https", "wss": - return tls.DialWithDialer(d, "tcp", hostPort, &tls.Config{ServerName: dest.Hostname()}) - default: - return d.Dial("tcp", hostPort) - } -} - -func rawHTTPError(conn net.Conn, code int, text string) error { - reason := http.StatusText(code) - if reason == "" { - reason = "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) - return err } // MihomoMetaJSON returns a JSON body for GET .../mihomo/meta (display URL without credentials). diff --git a/internal/proxy/mihomo_test.go b/internal/proxy/mihomo_test.go index b76f284..554b10d 100644 --- a/internal/proxy/mihomo_test.go +++ b/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 -}