// 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(
"&", "&", "<", "<", ">", ">", `"`, """, "'", "'",
).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 = `
登录 · ModelRouter
`
// 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)
}
}