telemt-api/internal/proxy/mihomo.go

package proxy

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// 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.
func NewMihomoForward(
	target *url.URL,
	stripPrefix string,
	auth string,
	rt http.RoundTripper,
	errHandler func(http.ResponseWriter, *http.Request, error),
) http.Handler {
	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) {
			wsH.ServeHTTP(w, r)
			return
		}
		httpH.ServeHTTP(w, r)
	})
}

func isWebSocketUpgrade(r *http.Request) bool {
	if r == nil {
		return false
	}
	if headerContainsToken(r.Header, "Upgrade", "websocket") &&
		headerContainsToken(r.Header, "Connection", "upgrade") {
		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
			}
		}
	}
	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 newMihomoWSTunnel(
	target *url.URL,
	stripPrefix string,
	auth string,
	errHandler func(http.ResponseWriter, *http.Request, error),
) http.Handler {
	if errHandler == nil {
		errHandler = defaultForwardErrorHandler
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		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 not supported by ResponseWriter"))
			return
		}
		clientConn, clientBuf, err := hj.Hijack()
		if err != nil {
			errHandler(w, r, fmt.Errorf("websocket: hijack failed: %w", err))
			return
		}
		defer clientConn.Close()

		upConn, err := dialUpstream(dest)
		if err != nil {
			_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream connect failed")
			return
		}
		defer upConn.Close()

		reqBytes := buildWSUpgradeRequest(r, dest, auth)
		if _, err := upConn.Write(reqBytes); err != nil {
			_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream write failed")
			return
		}

		upBuf := bufio.NewReader(upConn)
		resp, err := http.ReadResponse(upBuf, nil)
		if err != nil {
			_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream response read failed")
			return
		}

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

// 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"
		}
	}
	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).
func MihomoMetaJSON(target *url.URL) []byte {
	display := strings.TrimSuffix(target.String(), "/")
	b, _ := json.Marshal(map[string]any{
		"ok":              true,
		"controller_base": display,
	})
	return b
}

// MihomoJSONError writes a JSON error for Mihomo routes when proxy is not configured.
func MihomoJSONError(w http.ResponseWriter, code, msg string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusNotFound)
	_ = json.NewEncoder(w).Encode(map[string]any{
		"ok":    false,
		"error": map[string]string{"code": code, "message": msg},
	})
}