telemt-api/internal/server/gateway.go

package server

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/telemt/telemt-api/internal/aggregate"
	"github.com/telemt/telemt-api/internal/config"
	"github.com/telemt/telemt-api/internal/geoip"
	"github.com/telemt/telemt-api/internal/proxy"
	"github.com/telemt/telemt-api/internal/webui"
)

// Gateway serves health, metrics, and proxied API routes.
type Gateway struct {
	parsed        *config.Parsed
	proxies       map[string]http.Handler
	mihomo        map[string]http.Handler
	agg           *aggregate.Handler
	geo           *geoip.Service
	log           *slog.Logger
	transport     *http.Transport
	promHandler   http.Handler
	corsAllowed   []string
	webUI         http.Handler
	radarStatusesURL string
	radarPingFrom    string
	radarHTTPClient  *http.Client
}

func writeBadGatewayJSON(w http.ResponseWriter, expose bool, code, message string, upstreamErr error) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(http.StatusBadGateway)
	errObj := map[string]any{"code": code, "message": message}
	if expose && upstreamErr != nil {
		errObj["detail"] = upstreamErr.Error()
	}
	_ = json.NewEncoder(w).Encode(map[string]any{
		"ok":    false,
		"error": errObj,
	})
}

// NewGateway builds handlers and reverse proxies from parsed config.
func NewGateway(p *config.Parsed, log *slog.Logger, geo *geoip.Service) (*Gateway, error) {
	t := proxy.DirectTransport()
	t.MaxIdleConns = 64
	t.IdleConnTimeout = 90 * time.Second
	t.TLSHandshakeTimeout = 10 * time.Second
	t.ExpectContinueTimeout = 1 * time.Second
	t.DialContext = (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext
	t.ResponseHeaderTimeout = 120 * time.Second
	g := &Gateway{
		parsed:      p,
		proxies:     make(map[string]http.Handler),
		mihomo:      make(map[string]http.Handler),
		geo:         geo,
		log:         log,
		transport:   t,
		promHandler: promhttp.Handler(),
	}
	exposeErr := p.Config.ExposeUpstreamErrors
	for i := range p.Config.Servers {
		s := &p.Config.Servers[i]
		u, err := url.Parse(s.BaseURL)
		if err != nil {
			return nil, err
		}
		auth := p.AuthByAlias[s.Alias]
		strip := "/api/" + s.Alias
		g.proxies[s.Alias] = proxy.NewAliasForward(u, strip, s.PathPrefix, auth, t, func(w http.ResponseWriter, r *http.Request, err error) {
			log.Error("upstream error", "alias", s.Alias, "err", err)
			writeBadGatewayJSON(w, exposeErr, "bad_gateway", "upstream unreachable", err)
		})
	}
	for alias, m := range p.MihomoByAlias {
		a := alias
		strip := "/api/" + a + "/mihomo"
		g.mihomo[a] = proxy.NewMihomoForward(m.Base, strip, m.Auth, t, func(w http.ResponseWriter, r *http.Request, err error) {
			log.Error("mihomo upstream error", "alias", a, "err", err)
			writeBadGatewayJSON(w, exposeErr, "bad_gateway", "mihomo upstream unreachable", err)
		})
	}
	var aggCacheTTL time.Duration
	if p.Config.Aggregate != nil && p.Config.Aggregate.CacheTTLMs > 0 {
		aggCacheTTL = time.Duration(p.Config.Aggregate.CacheTTLMs) * time.Millisecond
	}
	g.corsAllowed = append([]string(nil), p.Config.CorsAllowedOrigins...)
	g.agg = aggregate.NewHandler(p, &http.Client{Transport: t}, geo, aggCacheTTL)
	g.webUI = webui.Handler()

	statusesURL := defaultRadarStatusesURL
	radarTo := 30 * time.Second
	if p.Config.Radar != nil {
		if u := strings.TrimSpace(p.Config.Radar.StatusesURL); u != "" {
			statusesURL = u
		}
		if p.Config.Radar.HTTPTimeoutMs > 0 {
			radarTo = time.Duration(p.Config.Radar.HTTPTimeoutMs) * time.Millisecond
		}
		g.radarPingFrom = strings.TrimSpace(p.Config.Radar.PingFrom)
	}
	g.radarStatusesURL = statusesURL
	g.radarHTTPClient = &http.Client{Transport: t, Timeout: radarTo}

	return g, nil
}

// Handler returns the root HTTP handler with middleware.
func (g *Gateway) Handler() http.Handler {
	var h http.Handler = http.HandlerFunc(g.serve)
	h = g.withCORS(h)
	h = g.withWhitelist(h)
	h = g.withAccessLog(h)
	h = g.withMetrics(h)
	return h
}

func (g *Gateway) withCORS(next http.Handler) http.Handler {
	if len(g.corsAllowed) == 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Vary", "Origin")
		origin := r.Header.Get("Origin")
		ok, allowOrigin := corsMatch(g.corsAllowed, origin)
		if ok {
			w.Header().Set("Access-Control-Allow-Origin", allowOrigin)
			w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, PATCH, DELETE, OPTIONS")
			w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Request-Id")
			w.Header().Set("Access-Control-Max-Age", "86400")
		}
		if r.Method == http.MethodOptions {
			if ok {
				w.WriteHeader(http.StatusNoContent)
				return
			}
		}
		next.ServeHTTP(w, r)
	})
}

func corsMatch(allowed []string, origin string) (ok bool, allowOrigin string) {
	if origin == "" {
		return false, ""
	}
	for _, a := range allowed {
		a = strings.TrimSpace(a)
		if a == "" {
			continue
		}
		if a == "*" {
			return true, "*"
		}
		if strings.EqualFold(a, origin) {
			return true, origin
		}
	}
	return false, ""
}

func (g *Gateway) withWhitelist(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/health" {
			next.ServeHTTP(w, r)
			return
		}
		ip := ClientIP(r, g.parsed.Trusted)
		if !Allowed(ip, g.parsed.Config.AllowAll, g.parsed.Whitelist) {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusForbidden)
			_ = json.NewEncoder(w).Encode(map[string]any{
				"ok":    false,
				"error": map[string]string{"code": "forbidden", "message": "source address not allowed"},
			})
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (g *Gateway) withAccessLog(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rid := r.Header.Get("X-Request-Id")
		if rid == "" {
			rid = randomID()
			r.Header.Set("X-Request-Id", rid)
		}
		w.Header().Set("X-Request-Id", rid)
		start := time.Now()
		lw := &statusWriter{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(lw, r)
		g.log.Info("request",
			"request_id", rid,
			"method", r.Method,
			"path", r.URL.Path,
			"status", lw.status,
			"duration_ms", time.Since(start).Milliseconds(),
			"remote", r.RemoteAddr,
		)
	})
}

func (g *Gateway) withMetrics(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/health" {
			next.ServeHTTP(w, r)
			return
		}
		httpInFlight.Inc()
		start := time.Now()
		alias := routeAlias(r.URL.Path)
		endpoint := routeEndpoint(r.URL.Path)
		lw := &statusWriter{ResponseWriter: w, status: http.StatusOK}
		defer observeRequest(r.Method, alias, endpoint, lw.status, start)
		next.ServeHTTP(lw, r)
	})
}

func routeAlias(path string) string {
	const pfx = "/api/"
	if !strings.HasPrefix(path, pfx) {
		if path == "/metrics" {
			return "metrics"
		}
		return "_"
	}
	rest := strings.TrimPrefix(path, pfx)
	if rest == "" {
		return "_"
	}
	i := strings.IndexByte(rest, '/')
	if i < 0 {
		return rest
	}
	return rest[:i]
}

func routeEndpoint(path string) string {
	if path == "" || path == "/" {
		return "ui_root"
	}
	if path == "/health" || path == "/metrics" {
		return strings.TrimPrefix(path, "/")
	}
	if strings.HasPrefix(path, "/api/agg/") {
		return strings.TrimPrefix(path, "/api/agg/")
	}
	if path == "/api/agg" {
		return "agg"
	}
	if strings.HasPrefix(path, "/api/radar/") {
		rest := strings.TrimPrefix(path, "/api/radar/")
		if rest == "" {
			return "radar"
		}
		return "radar_" + strings.ReplaceAll(rest, "/", "_")
	}
	if strings.HasPrefix(path, "/api/live/events") {
		return "live_events"
	}
	if strings.HasPrefix(path, "/api/") {
		rest := strings.TrimPrefix(path, "/api/")
		i := strings.IndexByte(rest, '/')
		if i < 0 {
			return "proxy_root"
		}
		sub := rest[i+1:]
		if strings.HasPrefix(sub, "mihomo/") {
			return "mihomo_" + sub[len("mihomo/"):]
		}
		return "proxy_" + sub
	}
	return "ui"
}

func (g *Gateway) serve(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/health":
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		_ = json.NewEncoder(w).Encode(map[string]any{"status": "ok"})
		return
	case "/metrics":
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		g.promHandler.ServeHTTP(w, r)
		return
	}
	// So /api//mtg/health and /api/../api/agg/... route like /api/mtg/health and /api/agg/...
	if strings.HasPrefix(r.URL.Path, "/api") {
		proxy.NormalizeRequestURLPath(r)
	}
	const prefix = "/api/"
	if r.Method == http.MethodPost && r.URL.Path == "/api/agg/bad-connections/reset" {
		g.agg.HandleResetBadConnectionsBaseline(w, r)
		return
	}
	if r.URL.Path == "/api/agg" || strings.HasPrefix(r.URL.Path, "/api/agg/") {
		g.agg.ServeHTTP(w, r)
		return
	}
	if r.URL.Path == "/api/live/events" {
		g.serveLiveEvents(w, r)
		return
	}
	if r.URL.Path == "/api/radar/statuses" || r.URL.Path == "/api/radar/ping-dc" {
		g.serveRadar(w, r)
		return
	}
	if !strings.HasPrefix(r.URL.Path, prefix) {
		g.webUI.ServeHTTP(w, r)
		return
	}
	trim := strings.TrimPrefix(r.URL.Path, prefix)
	if trim == "" {
		http.NotFound(w, r)
		return
	}
	var alias string
	if i := strings.IndexByte(trim, '/'); i >= 0 {
		alias = trim[:i]
	} else {
		alias = trim
	}
	if alias == "" {
		http.NotFound(w, r)
		return
	}
	_, ok := g.proxies[alias]
	if !ok {
		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": "not_found", "message": "unknown alias"},
		})
		return
	}
	mihomoPfx := "/api/" + alias + "/mihomo"
	if strings.HasPrefix(r.URL.Path, mihomoPfx) {
		mh, have := g.mihomo[alias]
		if !have {
			proxy.MihomoJSONError(w, "mihomo_not_configured", "mihomo is not configured for this server")
			return
		}
		if r.URL.Path == mihomoPfx+"/meta" && r.Method == http.MethodGet {
			mu := g.parsed.MihomoByAlias[alias]
			if mu == nil {
				proxy.MihomoJSONError(w, "mihomo_not_configured", "mihomo is not configured for this server")
				return
			}
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write(proxy.MihomoMetaJSON(mu.Base))
			return
		}
		mh.ServeHTTP(w, r)
		return
	}
	g.proxies[alias].ServeHTTP(w, r)
}

func (g *Gateway) serveLiveEvents(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "stream unsupported", http.StatusInternalServerError)
		return
	}
	aliases, err := g.agg.ResolveAliasesForLive(r)
	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(http.StatusBadRequest)
		_ = json.NewEncoder(w).Encode(map[string]any{
			"ok":    false,
			"error": map[string]string{"code": "bad_request", "message": err.Error()},
		})
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	push := func() {
		ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
		defer cancel()
		payload := aggregate.BuildLiveEnvelope(ctx, g.agg, aliases)
		b, _ := json.Marshal(payload)
		_, _ = fmt.Fprintf(w, "event: snapshot\n")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", string(b))
		flusher.Flush()
	}
	push()
	for {
		select {
		case <-r.Context().Done():
			return
		case <-ticker.C:
			push()
		}
	}
}

type statusWriter struct {
	http.ResponseWriter
	status int
}

func (s *statusWriter) WriteHeader(code int) {
	s.status = code
	s.ResponseWriter.WriteHeader(code)
}

// Flush keeps streaming compatibility (SSE/EventSource) through middleware wrappers.
func (s *statusWriter) Flush() {
	if f, ok := s.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// Hijack preserves WebSocket upgrade support through middleware wrappers.
func (s *statusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if h, ok := s.ResponseWriter.(http.Hijacker); ok {
		return h.Hijack()
	}
	return nil, nil, http.ErrNotSupported
}

// ReadFrom keeps io.Copy fast-path behavior when the underlying writer supports it.
func (s *statusWriter) ReadFrom(r io.Reader) (int64, error) {
	if rf, ok := s.ResponseWriter.(io.ReaderFrom); ok {
		return rf.ReadFrom(r)
	}
	return io.Copy(s.ResponseWriter, r)
}

// Unwrap lets net/http.ResponseController reach the original writer.
func (s *statusWriter) Unwrap() http.ResponseWriter {
	return s.ResponseWriter
}

// Shutdown idle connections on the shared transport.
func (g *Gateway) Shutdown(ctx context.Context) error {
	g.transport.CloseIdleConnections()
	if g.geo != nil {
		_ = g.geo.Close()
	}
	return nil
}