telemt-api/internal/proxy/mihomo.go

package proxy

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

// NewMihomoForward proxies /api/{alias}/mihomo/... to Mihomo external-controller.
//
// 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,
	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) {
		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
		}
		httpH.ServeHTTP(w, r)
	})
}

// useMihomoWebSocketTunnel is true when the client is doing a WS handshake.
// For /traffic, /memory, /connections 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 websocket.IsWebSocketUpgrade(r) {
		return true
	}
	if r.Method != http.MethodGet {
		return false
	}
	switch rest {
	case "traffic", "memory", "connections":
		return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != ""
	default:
		return false
	}
}

// applyMihomoWSDialAuth prepares upstream WS URL and Dial headers.
// Mihomo: /connections WebSocket expects ?token=<secret> (raw secret, no "Bearer " prefix);
// /traffic, /memory use Authorization: Bearer … like REST.
func applyMihomoWSDialAuth(rest string, dest *url.URL, auth string) http.Header {
	hdr := make(http.Header)
	a := strings.TrimSpace(auth)
	if a == "" || dest == nil {
		return hdr
	}
	if rest == "connections" {
		token := a
		if len(token) > 7 && strings.EqualFold(token[:7], "Bearer ") {
			token = strings.TrimSpace(token[7:])
		}
		q := dest.Query()
		q.Set("token", token)
		dest.RawQuery = q.Encode()
		return hdr
	}
	hdr.Set("Authorization", a)
	return hdr
}

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

		dialHdr := applyMihomoWSDialAuth(rest, dest, auth)
		wsURL, err := mihomoWebSocketURL(dest)
		if err != nil {
			errHandler(w, r, err)
			return
		}

		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 {
			errHandler(w, r, fmt.Errorf("mihomo ws dial: %w", err))
			return
		}
		defer upConn.Close()

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

		upgrader := websocket.Upgrader{
			HandshakeTimeout: 15 * time.Second,
			CheckOrigin:      func(*http.Request) bool { return true },
		}
		clientConn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer clientConn.Close()

		errCh := make(chan error, 2)
		go relayWSFrames(upConn, clientConn, errCh)
		go relayWSFrames(clientConn, upConn, errCh)
		<-errCh
	})
}

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

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