Implement WebSocket proxy handling in Mihomo

- Added a new function to handle WebSocket connections, including hijacking the client connection and establishing a connection to the upstream WebSocket server.
- Enhanced error handling for connection issues and improved request cloning for WebSocket upgrades.
- Introduced utility functions for dialing WebSocket upstream and writing raw HTTP errors, ensuring robust communication and error reporting.
- Refactored the existing proxy logic to accommodate the new WebSocket handling, improving overall functionality and reliability.

internal/proxy/mihomo.go

@@ -1,11 +1,17 @@
 package proxy
 
 import (
+	"bufio"
+	"crypto/tls"
 	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net"
 	"net/http"
-	"net/http/httputil"
 	"net/url"
 	"strings"
+	"time"
 )
 
 // NewMihomoForward proxies /api/{alias}/mihomo/... to Mihomo external-controller.
@@ -67,40 +73,129 @@ func newMihomoWebSocketReverseProxy(
 	if errHandler == nil {
 		errHandler = defaultForwardErrorHandler
 	}
-	if rt == nil {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		rt = http.DefaultTransport
+		NormalizeRequestURLPath(r)
+		p := r.URL.Path
+		if !strings.HasPrefix(p, stripPrefix) {
+			errHandler(w, r, fmt.Errorf("path %q: missing strip prefix %q", p, stripPrefix))
+			return
+		}
+		rest := strings.TrimPrefix(strings.TrimPrefix(p, stripPrefix), "/")
+		dest := JoinPathPrefix(target, "/", rest)
+		du := *dest
+		du.RawQuery = r.URL.RawQuery
+		dest = &du
+
+		hj, ok := w.(http.Hijacker)
+		if !ok {
+			errHandler(w, r, fmt.Errorf("websocket hijack unsupported"))
+			return
+		}
+		clientConn, clientRW, err := hj.Hijack()
+		if err != nil {
+			errHandler(w, r, fmt.Errorf("hijack client conn: %w", err))
+			return
+		}
+		defer clientConn.Close()
+
+		upConn, err := dialWebSocketUpstream(dest)
+		if err != nil {
+			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
+			return
+		}
+		defer upConn.Close()
+
+		outReq := r.Clone(r.Context())
+		outReq.URL = dest
+		outReq.Host = dest.Host
+		outReq.RequestURI = ""
+		outReq.Proto = "HTTP/1.1"
+		outReq.ProtoMajor = 1
+		outReq.ProtoMinor = 1
+		outReq.Header = cloneHeader(r.Header)
+		outReq.Header.Del("Authorization")
+		if auth != "" {
+			outReq.Header.Set("Authorization", auth)
+		}
+		if err := outReq.Write(upConn); err != nil {
+			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
+			return
+		}
+
+		upBr := bufio.NewReader(upConn)
+		resp, err := http.ReadResponse(upBr, outReq)
+		if err != nil {
+			_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
+			return
+		}
+		if err := resp.Write(clientConn); err != nil {
+			return
+		}
+		if resp.StatusCode != http.StatusSwitchingProtocols {
+			return
+		}
+
+		clientSrc := io.MultiReader(clientRW, clientConn)
+		upSrc := io.MultiReader(upBr, upConn)
+		errCh := make(chan error, 2)
+		go proxyCopy(errCh, upConn, clientSrc)
+		go proxyCopy(errCh, clientConn, upSrc)
+		<-errCh
+	})
+}
+
+func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) {
+	if dest == nil {
+		return nil, fmt.Errorf("nil destination")
 	}
-	rp := &httputil.ReverseProxy{
+	hostPort := dest.Host
-		Rewrite: func(pr *httputil.ProxyRequest) {
+	if !strings.Contains(hostPort, ":") {
-			NormalizeRequestURLPath(pr.Out)
+		switch strings.ToLower(dest.Scheme) {
-			p := pr.Out.URL.Path
+		case "https", "wss":
-			if !strings.HasPrefix(p, stripPrefix) {
+			hostPort += ":443"
-				return
+		default:
-			}
+			hostPort += ":80"
-			rest := strings.TrimPrefix(strings.TrimPrefix(p, stripPrefix), "/")
+		}
-			dest := JoinPathPrefix(target, "/", rest)
-			du := *dest
-			du.RawQuery = pr.Out.URL.RawQuery
-			out := pr.Out
-			out.URL = &du
-			out.Header.Del("Host")
-			out.Host = du.Host
-			out.RequestURI = ""
-			out.Proto = "HTTP/1.1"
-			out.ProtoMajor = 1
-			out.ProtoMinor = 1
-			out.Header.Del("Authorization")
-			if auth != "" {
-				out.Header.Set("Authorization", auth)
-			}
-		},
-		Transport: rt,
-		FlushInterval: -1,
-		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
-			errHandler(w, r, err)
-		},
 	}
-	return rp
+	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 writeRawHTTPError(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
+}
+
+func proxyCopy(errCh chan<- error, dst io.Writer, src io.Reader) {
+	_, err := io.Copy(dst, src)
+	if err != nil && !isNetClosed(err) {
+		errCh <- err
+		return
+	}
+	errCh <- nil
+}
+
+func isNetClosed(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, net.ErrClosed) {
+		return true
+	}
+	msg := err.Error()
+	return strings.Contains(msg, "use of closed network connection")
 }
 
 // MihomoMetaJSON returns a JSON body for GET .../mihomo/meta (display URL without credentials).