ModelRouter/internal/gateway/server.go

// Package gateway exposes an OpenAI-compatible HTTP API over the provider
// registry: POST /v1/chat/completions (SDK + SSE), POST /v1/images/generations,
// GET /v1/models, protected by shared gateway API keys, plus a web UI and
// management API for adapters and sources.
package gateway

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"io/fs"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/core"
)

//go:embed ui/*
var uiFS embed.FS

// Gateway is the HTTP handler for the OpenAI-compatible endpoint + web UI.
type Gateway struct {
	core      *core.Core
	ui        http.Handler
	stats     *Stats
	probeMu   sync.Mutex
	lastProbe time.Time
	// loginGuard throttles /api/login brute-force attempts: per-IP failure
	// counters with exponential backoff, plus a global cap so a distributed
	// spray cannot outrun the per-IP window. Guarded by loginMu.
	loginMu    sync.Mutex
	loginFails map[string]*loginFail
	loginWin   int64 // unix sec of current global window start
	loginGlob  int   // failures in the global window
}

// loginFail tracks consecutive failed logins from one source IP.
type loginFail struct {
	count  int
	until  int64 // locked-out until this unix second (0 = not locked)
	lastAt int64
}

const (
	loginMaxFails     = 5       // failures before lockout kicks in
	loginLockBase     = 30      // first lockout: 30s
	loginLockCap      = 30 * 60 // repeated lockouts cap at 30min
	loginDecay        = 15 * 60 // counters decay after 15min of quiet
	loginGlobalMax    = 100     // global failures per minute across all IPs
	loginGlobalWindow = 60      // global window length (sec)
)

// loginAllow checks (and records) a login attempt for ip. It returns false
// when the attempt must be rejected: the IP is in exponential lockout or the
// global failure budget for the current window is exhausted.
func (g *Gateway) loginAllow(ip string, now int64) bool {
	g.loginMu.Lock()
	defer g.loginMu.Unlock()
	if g.loginFails == nil {
		g.loginFails = map[string]*loginFail{}
	}
	// global window rollover
	if now-g.loginWin >= loginGlobalWindow {
		g.loginWin = now
		g.loginGlob = 0
	}
	f := g.loginFails[ip]
	if f != nil && f.count > 0 && now-f.lastAt > loginDecay {
		delete(g.loginFails, ip) // quiet long enough: forgive and forget
		f = nil
	}
	if f != nil && f.until > now {
		return false // still locked out
	}
	if g.loginGlob >= loginGlobalMax {
		return false // global spray budget exhausted for this window
	}
	g.loginGlob++
	return true
}

// loginRecord notes the outcome of one attempt: failures escalate toward an
// exponential lockout; a success clears the IP's counter entirely.
func (g *Gateway) loginRecord(ip string, ok bool, now int64) {
	g.loginMu.Lock()
	defer g.loginMu.Unlock()
	if ok {
		if g.loginFails != nil {
			delete(g.loginFails, ip)
		}
		return
	}
	f := g.loginFails[ip]
	if f == nil {
		f = &loginFail{}
		g.loginFails[ip] = f
	}
	f.count++
	f.lastAt = now
	if f.count >= loginMaxFails {
		// lock duration doubles with each extra burst of failures, capped
		epoch := int64(f.count/loginMaxFails) - 1
		dur := int64(loginLockBase) << min(epoch, 6)
		if dur > loginLockCap {
			dur = loginLockCap
		}
		f.until = now + dur
	}
}

func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
	sub, err := fs.Sub(uiFS, "ui")
	if err != nil {
		return nil, err
	}
	st := NewStats(10000)
	if cfg := c.Config(); cfg != nil && cfg.RuntimeFile != "" {
		st.LoadAudit(cfg.RuntimeFile + ".audit.jsonl")
	}
	return &Gateway{
		core:  c,
		ui:    http.FileServer(http.FS(sub)),
		stats: st,
	}, nil
}

func (g *Gateway) Handler() http.Handler {
	// access log wraps all requests (auth, unauthenticated, login).
	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// login entry point + login/logout API are the only unauthenticated routes
		if r.URL.Path == "/login" || r.URL.Path == "/api/login" || r.URL.Path == "/api/logout" {
			g.routes(w, r)
			return
		}
		g.auth(http.HandlerFunc(g.routes)).ServeHTTP(w, r)
	})
	return g.logAccess(inner)
}

// logAccess wraps an http.Handler so every response is written to the audit
// file as an access event (method, path, status, latency, key).
func (g *Gateway) logAccess(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t0 := time.Now()
		sr := &statusRecorder{ResponseWriter: w}
		key := ""
		if h := r.Header.Get("Authorization"); h != "" {
			parts := strings.SplitN(h, " ", 2)
			if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
				key = parts[1]
			}
		}
		if key == "" {
			key = r.URL.Query().Get("api_key")
		}
		if key == "" {
			if c, err := r.Cookie("gw_key"); err == nil {
				key = c.Value
			}
		}
		next.ServeHTTP(sr, r)
		g.stats.AppendAudit("access", map[string]interface{}{
			"method": r.Method,
			"path":   r.URL.RequestURI(),
			"status": sr.code,
			"lat_ms": time.Since(t0).Milliseconds(),
			"key":    keyID(key),
		})
	})
}

type statusRecorder struct {
	http.ResponseWriter
	code int
}

func (s *statusRecorder) WriteHeader(code int) {
	if s.code == 0 {
		s.code = code
	}
	s.ResponseWriter.WriteHeader(code)
}

func (s *statusRecorder) Write(b []byte) (int, error) {
	if s.code == 0 {
		s.code = http.StatusOK
	}
	return s.ResponseWriter.Write(b)
}

// Flush forwards to the underlying writer so SSE handlers can stream
// incrementally through the access-log wrapper. Without this method the
// w.(http.Flusher) assertion inside the streaming chat handlers fails (the
// embedded ResponseWriter interface does not carry Flush into the method set),
// and every chunk stays buffered until the response ends.
func (s *statusRecorder) Flush() {
	if f, ok := s.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.URL.Path == "/v1/chat/completions":
		g.handleChat(w, r)
	case r.URL.Path == "/v1/images/generations":
		g.handleImage(w, r)
	case r.URL.Path == "/v1/models":
		g.handleModels(w, r)
	case r.URL.Path == "/api/adapters" || strings.HasPrefix(r.URL.Path, "/api/adapters/"):
		g.handleAdaptersAPI(w, r)
	case r.URL.Path == "/api/sources" || strings.HasPrefix(r.URL.Path, "/api/sources/"):
		g.handleSourcesAPI(w, r)
	case r.URL.Path == "/api/source_templates" || strings.HasPrefix(r.URL.Path, "/api/source_templates/"):
		g.handleSourceTemplatesAPI(w, r)
	case r.URL.Path == "/api/chat":
		g.handleChat(w, r)
	case r.URL.Path == "/api/status":
		g.handleStatusAPI(w, r)
	case r.URL.Path == "/api/status/reset":
		g.handleResetHealth(w, r)
	case r.URL.Path == "/api/stats/records":
		g.handleStatsRecordsAPI(w, r)
	case r.URL.Path == "/api/stats" || strings.HasPrefix(r.URL.Path, "/api/stats/"):
		g.handleStatsAPI(w, r)
	case r.URL.Path == "/api/keys" || strings.HasPrefix(r.URL.Path, "/api/keys/"):
		g.handleKeysAPI(w, r)
	case r.URL.Path == "/api/auto":
		g.handleAutoAPI(w, r)
	case r.URL.Path == "/login":
		g.handleLogin(w, r)
	case r.URL.Path == "/api/login":
		g.handleLoginAPI(w, r)
	case r.URL.Path == "/api/logout":
		g.handleLogoutAPI(w, r)
	default:
		g.serveUI(w, r)
	}
}

func (g *Gateway) serveUI(w http.ResponseWriter, r *http.Request) {
	// serve index.html directly for the root path (FileServer would 301 it)
	if r.URL.Path == "/" || r.URL.Path == "/ui" {
		data, err := uiFS.ReadFile("ui/index.html")
		if err != nil {
			http.Error(w, "ui missing", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		w.Write(data)
		return
	}
	g.ui.ServeHTTP(w, r)
}

func (g *Gateway) auth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := ""
		if h := r.Header.Get("Authorization"); h != "" {
			parts := strings.SplitN(h, " ", 2)
			if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
				key = parts[1]
			}
		}
		if key == "" {
			key = r.URL.Query().Get("api_key")
		}
		if key == "" {
			if c, err := r.Cookie("gw_key"); err == nil {
				key = c.Value
			}
		}
		rec, ok := g.core.FindKey(key)
		if !ok {
			if isAPIPath(r.URL.Path) {
				writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
				return
			}
			// browser navigation to UI pages -> login page
			http.Redirect(w, r, "/login?continue="+url.QueryEscape(r.URL.Path), http.StatusFound)
			return
		}
		next.ServeHTTP(w, r.WithContext(withAuth(r.Context(), key, rec.Role)))
	})
}

// authCtx carries the authenticated gateway key and its role.
type authCtx struct {
	key  string
	role string
}

type authCtxKeyT struct{}

func withAuth(ctx context.Context, key, role string) context.Context {
	return context.WithValue(ctx, authCtxKeyT{}, authCtx{key: key, role: role})
}

// reqKey returns the authenticated gateway key id (masked suffix for display).
func reqKey(ctx context.Context) string {
	if a, ok := ctx.Value(authCtxKeyT{}).(authCtx); ok {
		return a.key
	}
	return ""
}

// reqRole returns the authenticated key role ("admin" or "user").
func reqRole(ctx context.Context) string {
	if a, ok := ctx.Value(authCtxKeyT{}).(authCtx); ok {
		return a.role
	}
	return "user"
}

// keyID returns a short stable id for a gateway key (last 6 chars).
func keyID(k string) string {
	if len(k) <= 6 {
		return k
	}
	return "***" + k[len(k)-6:]
}

// isAPIPath reports whether the request targets a JSON API endpoint that
// should answer 401 instead of redirecting to the login page.
func isAPIPath(p string) bool {
	return strings.HasPrefix(p, "/v1/") || strings.HasPrefix(p, "/api/")
}

// handleLogoutAPI clears the session cookie and redirects to the login page.
func (g *Gateway) handleLogoutAPI(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "gw_key",
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
	})
	http.Redirect(w, r, "/login", http.StatusFound)
}

// handleLogin serves the login page (unauthenticated). Supports a
// ?continue= path to return to after a successful login and a ?lang=zh|en
// toggle for i18n.
func (g *Gateway) handleLogin(w http.ResponseWriter, r *http.Request) {
	continuePath := r.URL.Query().Get("continue")
	if continuePath == "" || !strings.HasPrefix(continuePath, "/") || strings.HasPrefix(continuePath, "//") {
		continuePath = "/"
	}
	lang := strings.ToLower(r.URL.Query().Get("lang"))
	if lang != "en" {
		lang = "zh"
	}
	// escape for embedding in HTML attribute and single-quoted JS string
	htmlCont := strings.NewReplacer(
		"&", "&amp;", "<", "&lt;", ">", "&gt;", `"`, "&quot;", "'", "&#39;",
	).Replace(continuePath)
	jsCont := strings.NewReplacer("\\", "\\\\", "'", "\\'", "\n", "\\n").Replace(continuePath)
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	fmt.Fprintf(w, loginPageHTML, htmlCont, jsCont, lang)
}

const loginPageHTML = `<!DOCTYPE html><html lang="zh"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>登录 · ModelRouter</title>
<style>
:root{--bg:#f4f6fb;--card:#ffffff;--line:#e2e6ef;--fg:#1c2333;--muted:#6b7390;--accent:#3f6ef5;--accent-h:#2f5ae0;--err:#e5484d}
*{box-sizing:border-box}
body{margin:0;font:14px/1.6 -apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,"PingFang SC","Microsoft YaHei",sans-serif;background:linear-gradient(160deg,#f4f6fb 0%%,#e8edf8 100%%);color:var(--fg);display:flex;align-items:center;justify-content:center;min-height:100vh;padding:20px}
.card{background:var(--card);border:1px solid var(--line);border-radius:14px;padding:36px 34px;width:340px;box-shadow:0 12px 40px rgba(31,45,90,.10)}
.brand{display:flex;align-items:center;gap:10px;margin-bottom:6px}
.logo{width:30px;height:30px;border-radius:8px;background:linear-gradient(135deg,#3f6ef5,#6a8ffb);display:flex;align-items:center;justify-content:center;color:#fff;font-weight:700;font-size:15px}
h1{font-size:18px;margin:0;letter-spacing:.2px}
.sub{color:var(--muted);font-size:13px;margin:4px 0 22px}
label{display:block;font-size:12px;color:var(--muted);margin-bottom:6px}
input{width:100%%;background:#fbfcfe;border:1px solid var(--line);border-radius:8px;padding:10px 12px;font:inherit;color:var(--fg);outline:none;transition:border .15s,box-shadow .15s}
input:focus{border-color:var(--accent);box-shadow:0 0 0 3px rgba(63,110,245,.15)}
button{margin-top:16px;width:100%%;background:var(--accent);color:#fff;border:0;border-radius:8px;padding:11px;font:inherit;font-weight:600;cursor:pointer;transition:background .15s}
button:hover{background:var(--accent-h)}
#msg{color:var(--err);margin-top:12px;min-height:18px;font-size:13px}
.foot{margin-top:18px;text-align:center;font-size:12px;color:var(--muted)}
.foot a{color:var(--muted);text-decoration:none;border-bottom:1px dashed var(--line)}
</style></head>
<body><div class="card">
<div class="brand"><div class="logo">M</div><div>
<h1 data-i="title">ModelRouter</h1>
<div class="sub" data-i="sub">统一 LLM 网关 · 登录</div></div></div>
<input id="key" type="password" data-i-ph="ph" placeholder="Gateway API Key" autocomplete="off">
<button onclick="login()" data-i="btn">登录</button>
<div id="msg"></div>
<div class="foot"><a data-i="lang" href="?lang=en&continue=%[1]s">English</a></div>
</div>
<script>
var LANG='%[3]s';
var CONT='%[2]s';
var I={zh:{title:'ModelRouter',sub:'统一 LLM 网关 · 登录',ph:'输入网关 API Key',btn:'登录',lang:'English'},
   en:{title:'ModelRouter',sub:'Unified LLM Gateway · Sign in',ph:'Enter gateway API key',btn:'Sign in',lang:'中文'}};
 function apply(){var t=I[LANG]||I.zh;document.querySelector('[data-i=title]').textContent=t.title;
 document.querySelector('[data-i=sub]').textContent=t.sub;
 var ph=document.querySelector('#key');if(ph)ph.placeholder=t.ph;
 document.querySelector('[data-i=btn]').textContent=t.btn;
 document.querySelector('[data-i=lang]').textContent=t.lang;
 document.documentElement.lang=LANG==='zh'?'zh':'en';}
document.querySelector('[data-i=lang]').onclick=function(e){e.preventDefault();
 location.href='/login?lang='+(LANG==='zh'?'en':'zh')+'&continue='+encodeURIComponent(CONT);};
async function login(){
 var key=document.getElementById('key').value.trim();
 if(!key){document.getElementById('msg').textContent=(LANG==='zh'?'请输入 Key':'Enter a key');return;}
 var r=await fetch('/api/login',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({key})});
 var j=await r.json().catch(function(){return{};});
 if(r.ok){location.href=CONT||'/';}
 else{document.getElementById('msg').textContent=(j.error&&j.error.message)||(LANG==='zh'?'登录失败':'Login failed');}
}
document.getElementById('key').addEventListener('keydown',function(e){if(e.key==='Enter')login();});
apply();
</script></body></html>`

// handleLoginAPI validates the gateway key and issues a session cookie,
// guarded by per-IP + global brute-force throttling (see loginGuard).
func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	now := time.Now().Unix()
	if !g.loginAllow(clientIP(r), now) {
		// Don't reveal whether the key was right: answer the same 401 the
		// failure path uses, but skip the expensive key lookup.
		writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
		return
	}
	var body struct {
		Key string `json:"key"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		g.loginRecord(clientIP(r), false, now)
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json")
		return
	}
	if _, ok := g.core.FindKey(body.Key); !ok {
		g.loginRecord(clientIP(r), false, now)
		writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
		return
	}
	g.loginRecord(clientIP(r), true, now)
	http.SetCookie(w, &http.Cookie{
		Name:     "gw_key",
		Value:    body.Key,
		Path:     "/",
		MaxAge:   86400 * 30,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})
	writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
}

// clientIP returns the caller's IP for rate limiting. Behind a trusted
// reverse proxy (nginx) the X-Real-IP header is preferred; otherwise the
// remote address is used with the port stripped.
func clientIP(r *http.Request) string {
	if x := r.Header.Get("X-Real-IP"); x != "" {
		return x
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	models := g.core.Registry().ModelList()
	if allow := g.allowedModels(r.Context()); allow != nil {
		models = intersectModels(models, allow)
	}
	type modelObj struct {
		ID     string `json:"id"`
		Object string `json:"object"`
	}
	objs := make([]modelObj, 0, len(models))
	for _, m := range models {
		objs = append(objs, modelObj{ID: m, Object: "model"})
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"object": "list",
		"data":   objs,
	})
}

// ensureProbe triggers a live source probe at most once every 30s.
// The probe runs asynchronously so a slow/stuck upstream never blocks the
// status response — reachability data is eventually-consistent and the UI
// polls /api/stats every 3s anyway, so the next refresh picks it up.
func (g *Gateway) ensureProbe(ctx context.Context) {
	g.probeMu.Lock()
	due := time.Since(g.lastProbe) > 30*time.Second
	if due {
		g.lastProbe = time.Now()
	}
	g.probeMu.Unlock()
	if !due {
		return
	}
	go func() {
		probeCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
		defer cancel()
		g.core.Registry().ProbeAll(probeCtx)
	}()
}

// handleResetHealth (admin) clears the per-source backoff state so a fixed
// upstream or an edited AUTO priority chain becomes schedulable immediately.
func (g *Gateway) handleResetHealth(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	g.core.ResetHealth()
	g.stats.AppendAudit("config", map[string]interface{}{"action": "reset_health", "key": keyID(reqKey(r.Context()))})
	writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
}

func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
	g.ensureProbe(r.Context())
	host := r.Host
	if host == "" {
		host = g.core.Listen()
	}
	if strings.HasPrefix(host, ":") {
		host = "127.0.0.1" + host
	}
	// gateway_keys: return only the current authenticated user's key so the
	// connection snippet on the home page always shows the correct key
	// (previously every user saw the first admin key in the list).
	myKey := reqKey(r.Context())
	ks := make([]config.GWKey, 0, len(g.core.ListKeys()))
	for _, k := range g.core.ListKeys() {
		if k.Key == myKey {
			ks = append(ks, config.GWKey{
				Key:  k.Key,
				Role: k.Role,
				Name: k.Name,
			})
			break
		}
	}
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	baseURL := g.core.PublicBaseURL()
	if baseURL == "" {
		baseURL = scheme + "://" + host + "/v1"
	}
	models := g.core.Registry().ModelList()
	if reqRole(r.Context()) != "admin" {
		models = g.scopedModelList(r.Context(), models)
	}
	resp := map[string]interface{}{
		"default_model": g.core.DefaultModel(),
		"models":        models,
		"base_url":      baseURL,
		"gateway_keys":  ks,
	}
	if reqRole(r.Context()) == "admin" {
		sts := g.core.Registry().Status()
		recent := g.stats.SourceRecent(300)
		avgs := g.stats.SourceAverages(300)
		for i := range sts {
			if v, ok := recent[sts[i].Name]; ok {
				sts[i].RecentOK = v[0]
				sts[i].RecentErr = v[1]
			}
			if a, ok := avgs[sts[i].Name]; ok {
				sts[i].AvgFirstByteMs = a.AvgFirstByteMs
				sts[i].AvgTokPerS = a.AvgTokPerS
			}
		}
		resp["sources"] = sts
		resp["adapters"] = g.core.ListAdapters()
		// Elastic Lua pool sizing: created/idle/in_use vs the ceiling, plus the
		// current grow/shrink steps, so the algorithm is inspectable in the UI
		// instead of being a black box.
		resp["adapter_pools"] = g.core.VM().PoolStats()
	}
	writeJSON(w, http.StatusOK, resp)
}

// scopedModelList filters the full model set to only those allowed by the
// request's gateway key (used for users with a restricted model scope).
func (g *Gateway) scopedModelList(ctx context.Context, full []string) []string {
	allow := g.allowedModels(ctx)
	if allow == nil {
		return full
	}
	out := make([]string, 0, len(allow))
	for _, m := range allow {
		if m.Model == "" {
			continue
		}
		out = append(out, m.Model)
	}
	return out
}

func writeError(w http.ResponseWriter, code int, errType, msg string) {
	writeJSON(w, code, map[string]interface{}{
		"error": map[string]interface{}{"type": errType, "message": msg},
	})
}

func writeJSON(w http.ResponseWriter, code int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		log.Printf("[gateway] write json: %v", err)
	}
}