Files
ModelRouter/internal/gateway/server.go
JianFeeeee 5b6d6fc90d fix(gateway): keep the record ring small now that aggregates scan everything
The gateway still built Stats with NewStats(10000), a leftover from when the
ring WAS the source of the dashboard's numbers. With aggregates now computed
from the full audit history, a 10000-entry ring only means 10000 resident Req
structs — the live instance came up at 64 MB instead of ~40 MB, putting back
most of the memory the on-demand paging work removed.

The ring's only jobs are the status page's 5-minute SourceRecent /
SourceAverages windows and the dashboard's first screen, both of which fit in
defaultRingSize (500). Older records are paged from disk.

TestGatewayRingStaysSmall pins it so the constant cannot drift back up.
2026-08-30 09:34:25 +08:00

640 lines
21 KiB
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
}
// The record ring is deliberately small: aggregates are built from the full
// audit history at startup (see Stats.LoadAudit), so the ring only has to
// cover the status page's 5-minute source windows and the dashboard's first
// screen. Everything older is paged from disk by /api/stats/records.
// A large ring here would put the ~25 MB of resident Req structs straight
// back, which is exactly what the on-demand paging removed.
st := NewStats(defaultRingSize)
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)
}
}