telemt-api/internal/proxy/forward.go

package proxy

import (
	"fmt"
	"io"
	"net/http"
	"net/textproto"
	"net/url"
	"strings"
)

// NewAliasForward proxies /api/{alias}/… to Telemt using http.NewRequest(fullURL)+RoundTrip,
// matching aggregate server-side calls. httputil.ReverseProxy can produce request lines that
// strict origin servers reject with 400; this path matches a working curl to base_url.
func NewAliasForward(
	target *url.URL,
	stripPrefix, pathPrefix, auth string,
	rt http.RoundTripper,
	errHandler func(http.ResponseWriter, *http.Request, error),
) http.Handler {
	return newAliasForward(target, stripPrefix, pathPrefix, auth, rt, errHandler, true)
}

// newAliasForward is the shared HTTP forwarder. remapUsers maps top-level "users" → "stats/users"
// for Telemt; Mihomo must pass false so /users is not rewritten.
func newAliasForward(
	target *url.URL,
	stripPrefix, pathPrefix, auth string,
	rt http.RoundTripper,
	errHandler func(http.ResponseWriter, *http.Request, error),
	remapUsers bool,
) http.Handler {
	if errHandler == nil {
		errHandler = defaultForwardErrorHandler
	}
	if rt == nil {
		rt = http.DefaultTransport
	}
	targetQuery := target.RawQuery

	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(p, stripPrefix)
		rest = strings.TrimPrefix(rest, "/")
		if remapUsers && (r.Method == http.MethodGet || r.Method == http.MethodHead) && rest == "users" {
			rest = "stats/users"
		}

		outURL := JoinPathPrefix(target, pathPrefix, rest)
		u := *outURL
		if targetQuery == "" || r.URL.RawQuery == "" {
			u.RawQuery = targetQuery + r.URL.RawQuery
		} else {
			u.RawQuery = targetQuery + "&" + r.URL.RawQuery
		}
		outURL = &u

		outReq, err := http.NewRequestWithContext(r.Context(), r.Method, outURL.String(), r.Body)
		if err != nil {
			errHandler(w, r, err)
			return
		}
		if r.ContentLength >= 0 {
			outReq.ContentLength = r.ContentLength
		}
		outReq.Header = cloneHeader(r.Header)
		removeConnectionHeaders(outReq.Header)
		// Do not forward the client's Host (e.g. mtg.ivx.su:8888). Upstream must see the
		// authority from base_url (e.g. 172.20.0.3:9091); mismatch often yields 400 from strict stacks.
		outReq.Header.Del("Host")
		outReq.Host = outURL.Host
		if auth != "" {
			outReq.Header.Set("Authorization", auth)
		}

		resp, err := rt.RoundTrip(outReq)
		if err != nil {
			errHandler(w, r, err)
			return
		}
		defer resp.Body.Close()

		removeConnectionHeaders(resp.Header)
		for k, vv := range resp.Header {
			for _, v := range vv {
				w.Header().Add(k, v)
			}
		}
		w.WriteHeader(resp.StatusCode)
		_, _ = io.Copy(w, resp.Body)
	})
}

func defaultForwardErrorHandler(w http.ResponseWriter, _ *http.Request, _ error) {
	http.Error(w, "bad gateway", http.StatusBadGateway)
}

func cloneHeader(h http.Header) http.Header {
	h2 := make(http.Header, len(h))
	for k, vv := range h {
		cp := make([]string, len(vv))
		copy(cp, vv)
		h2[k] = cp
	}
	return h2
}

// removeConnectionHeaders mirrors net/http/httputil.ReverseProxy hop-by-hop handling.
func removeConnectionHeaders(h http.Header) {
	if v := h.Get("Connection"); v != "" {
		for _, f := range strings.Split(v, ",") {
			if f = textproto.TrimString(f); f != "" {
				h.Del(f)
			}
		}
	}
	for _, k := range hopHeaders {
		h.Del(k)
	}
}

var hopHeaders = []string{
	"Connection",
	"Proxy-Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailer",
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}