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

@@ -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

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).

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