328 lines
11 KiB
Go
328 lines
11 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"crypto/subtle"
|
|
"database/sql"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"io/fs"
|
|
_ "modernc.org/sqlite"
|
|
"net"
|
|
"net/http"
|
|
"net/netip"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type Server struct {
|
|
logRetentionDays atomic.Int64
|
|
trustedProxies []netip.Prefix
|
|
limiter *requestLimiter
|
|
db *sql.DB
|
|
password [32]byte
|
|
assets fs.FS
|
|
mu sync.Mutex
|
|
sessions map[string]time.Time
|
|
attempts map[string][]time.Time
|
|
authFailures map[string][]time.Time
|
|
authFailuresSeen time.Time
|
|
loginFailures []time.Time
|
|
client *http.Client
|
|
publicOrigin string
|
|
lifecycleMu sync.Mutex
|
|
closing bool
|
|
active sync.WaitGroup
|
|
closeOnce sync.Once
|
|
closeErr error
|
|
logMu sync.RWMutex
|
|
logsClosed bool
|
|
logQueue chan logEntry
|
|
logDone chan struct{}
|
|
logErr error
|
|
}
|
|
|
|
const maxBody = 16 << 20
|
|
|
|
func New(p, password string, a fs.FS) (*Server, error) {
|
|
if password == "" {
|
|
return nil, errors.New("empty admin password")
|
|
}
|
|
if p != ":memory:" {
|
|
if e := os.MkdirAll(filepath.Dir(p), 0700); e != nil {
|
|
return nil, e
|
|
}
|
|
}
|
|
d, e := sql.Open("sqlite", p)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
d.SetMaxOpenConns(1)
|
|
schema := `PRAGMA busy_timeout=5000; CREATE TABLE IF NOT EXISTS channels(id TEXT PRIMARY KEY,name TEXT NOT NULL UNIQUE,base_url TEXT NOT NULL,protocol TEXT NOT NULL,enabled INTEGER NOT NULL,passthrough INTEGER NOT NULL,models TEXT NOT NULL,api_key TEXT NOT NULL); CREATE TABLE IF NOT EXISTS keys(id TEXT PRIMARY KEY,name TEXT NOT NULL,prefix TEXT NOT NULL,hash TEXT NOT NULL UNIQUE,enabled INTEGER NOT NULL,all_channels INTEGER NOT NULL,channels TEXT NOT NULL,expires_at TEXT NOT NULL,created_at TEXT NOT NULL,revoked INTEGER NOT NULL DEFAULT 0); CREATE TABLE IF NOT EXISTS logs(id INTEGER PRIMARY KEY AUTOINCREMENT,created_at TEXT NOT NULL,key_id TEXT NOT NULL,key_name TEXT NOT NULL,channel TEXT NOT NULL,request_model TEXT NOT NULL,actual_model TEXT NOT NULL,downstream TEXT NOT NULL,upstream TEXT NOT NULL,status INTEGER NOT NULL,duration_ms INTEGER NOT NULL,stream INTEGER NOT NULL,input_tokens INTEGER,output_tokens INTEGER,cache_read_tokens INTEGER,cache_write_tokens INTEGER,usage_known INTEGER NOT NULL,error TEXT NOT NULL,raw_usage TEXT NOT NULL); CREATE INDEX IF NOT EXISTS logs_time ON logs(created_at); CREATE INDEX IF NOT EXISTS logs_key ON logs(key_id);`
|
|
if _, e = d.Exec(schema); e != nil {
|
|
d.Close()
|
|
return nil, e
|
|
}
|
|
if p != ":memory:" {
|
|
if e = os.Chmod(p, 0600); e != nil {
|
|
d.Close()
|
|
return nil, e
|
|
}
|
|
}
|
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
|
// Upstreams must not be able to turn compressed streams into proxy-side
|
|
// decompression work; the router asks for identity encoding.
|
|
tr.DisableCompression = true
|
|
tr.ResponseHeaderTimeout = 60 * time.Second
|
|
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
|
c, e := (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, addr)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
return &idleConn{Conn: c}, nil
|
|
}
|
|
s := &Server{db: d, password: sha256.Sum256([]byte(password)), assets: a, sessions: map[string]time.Time{}, attempts: map[string][]time.Time{}, authFailures: map[string][]time.Time{}, client: &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, logQueue: make(chan logEntry, 1024), logDone: make(chan struct{})}
|
|
go s.writeLogs()
|
|
s.limiter = newRequestLimiter(120, 8, 64)
|
|
return s, nil
|
|
}
|
|
func (s *Server) Close() error {
|
|
s.closeOnce.Do(func() {
|
|
s.lifecycleMu.Lock()
|
|
s.closing = true
|
|
s.lifecycleMu.Unlock()
|
|
s.active.Wait()
|
|
s.logMu.Lock()
|
|
s.logsClosed = true
|
|
close(s.logQueue)
|
|
s.logMu.Unlock()
|
|
<-s.logDone
|
|
s.client.CloseIdleConnections()
|
|
s.closeErr = errors.Join(s.logErr, s.db.Close())
|
|
})
|
|
return s.closeErr
|
|
}
|
|
func writeJSON(w http.ResponseWriter, v any) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(v)
|
|
}
|
|
func errJSON(w http.ResponseWriter, n int, m string) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(n)
|
|
json.NewEncoder(w).Encode(map[string]string{"error": m})
|
|
}
|
|
func decode(w http.ResponseWriter, r *http.Request, v any) error {
|
|
limit := int64(maxBody)
|
|
if r.URL.Path == "/admin/api/login" {
|
|
limit = 16 << 10
|
|
}
|
|
return readBody(w, r, limit, bodyReadTimeout, func() error {
|
|
d := json.NewDecoder(r.Body)
|
|
if e := d.Decode(v); e != nil {
|
|
return e
|
|
}
|
|
var extra any
|
|
if d.Decode(&extra) != io.EOF {
|
|
return errors.New("trailing JSON")
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
func randomID() string {
|
|
b := make([]byte, 32)
|
|
if _, e := rand.Read(b); e != nil {
|
|
panic(e)
|
|
}
|
|
return hex.EncodeToString(b)
|
|
}
|
|
func hash(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
|
|
func (s *Server) session(r *http.Request) bool {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
c, e := r.Cookie("ai_router_session")
|
|
if e != nil {
|
|
return false
|
|
}
|
|
t, ok := s.sessions[c.Value]
|
|
return ok && time.Now().Before(t)
|
|
}
|
|
func (s *Server) guard(r *http.Request) bool {
|
|
if r.Header.Get("X-Requested-With") != "ai-router" {
|
|
return false
|
|
}
|
|
o, e := url.Parse(r.Header.Get("Origin"))
|
|
if s.publicOrigin != "" {
|
|
// Origins are scheme/host case-insensitive; a mixed-case deployment
|
|
// value must not lock the admin out.
|
|
return e == nil && strings.EqualFold(r.Header.Get("Origin"), s.publicOrigin)
|
|
}
|
|
scheme := "http"
|
|
if r.TLS != nil {
|
|
scheme = "https"
|
|
}
|
|
return e == nil && o.Scheme == scheme && strings.EqualFold(o.Host, r.Host) && o.User == nil && o.RawQuery == "" && o.Fragment == "" && o.Path == ""
|
|
}
|
|
func (s *Server) secureCookie() bool {
|
|
return len(s.publicOrigin) >= 8 && strings.EqualFold(s.publicOrigin[:8], "https://")
|
|
}
|
|
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
w = &deadlineWriter{ResponseWriter: w}
|
|
_ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(bodyReadTimeout))
|
|
s.lifecycleMu.Lock()
|
|
if s.closing {
|
|
s.lifecycleMu.Unlock()
|
|
errJSON(w, http.StatusServiceUnavailable, "server shutting down")
|
|
return
|
|
}
|
|
s.active.Add(1)
|
|
s.lifecycleMu.Unlock()
|
|
defer s.active.Done()
|
|
w.Header().Set("X-Content-Type-Options", "nosniff")
|
|
w.Header().Set("Referrer-Policy", "no-referrer")
|
|
w.Header().Set("X-Frame-Options", "DENY")
|
|
w.Header().Set("Content-Security-Policy", "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'none'; connect-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:")
|
|
if strings.HasPrefix(r.URL.Path, "/admin/api/") || r.URL.Path == "/admin/api" {
|
|
w.Header().Set("Cache-Control", "no-store")
|
|
s.admin(w, r)
|
|
return
|
|
}
|
|
if strings.HasPrefix(r.URL.Path, "/v1/") {
|
|
s.proxy(w, r)
|
|
return
|
|
}
|
|
if strings.HasPrefix(r.URL.Path, "/api/") {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
if s.assets == nil || (r.Method != "GET" && r.Method != "HEAD") {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
p := strings.TrimPrefix(r.URL.Path, "/")
|
|
if p == "" {
|
|
p = "index.html"
|
|
}
|
|
if st, e := fs.Stat(s.assets, p); e != nil || st.IsDir() {
|
|
p = "index.html"
|
|
}
|
|
http.ServeFileFS(w, r, s.assets, p)
|
|
}
|
|
func (s *Server) admin(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != "GET" && !s.guard(r) {
|
|
errJSON(w, 403, "forbidden")
|
|
return
|
|
}
|
|
if r.URL.Path == "/admin/api/login" && r.Method == "POST" {
|
|
s.login(w, r)
|
|
return
|
|
}
|
|
if r.URL.Path == "/admin/api/session" && r.Method == "GET" {
|
|
writeJSON(w, map[string]bool{"authenticated": s.session(r)})
|
|
return
|
|
}
|
|
if !s.session(r) {
|
|
errJSON(w, 401, "unauthorized")
|
|
return
|
|
}
|
|
if r.URL.Path == "/admin/api/logout" && r.Method == "POST" {
|
|
c, _ := r.Cookie("ai_router_session")
|
|
s.mu.Lock()
|
|
delete(s.sessions, c.Value)
|
|
s.mu.Unlock()
|
|
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", MaxAge: -1, Path: "/", HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
|
writeJSON(w, map[string]bool{"authenticated": false})
|
|
return
|
|
}
|
|
switch {
|
|
case r.URL.Path == "/admin/api/models/test":
|
|
s.testModel(w, r)
|
|
case r.URL.Path == "/admin/api/channels":
|
|
s.channels(w, r)
|
|
case r.URL.Path == "/admin/api/channels/discover-models":
|
|
s.discoverModels(w, r)
|
|
case strings.HasPrefix(r.URL.Path, "/admin/api/channels/"):
|
|
s.channelID(w, r)
|
|
case r.URL.Path == "/admin/api/keys":
|
|
s.keys(w, r)
|
|
case strings.HasPrefix(r.URL.Path, "/admin/api/keys/"):
|
|
s.keyID(w, r)
|
|
case r.URL.Path == "/admin/api/logs" && r.Method == "GET":
|
|
s.logs(w, r)
|
|
case r.URL.Path == "/admin/api/stats" && r.Method == "GET":
|
|
s.stats(w, r)
|
|
default:
|
|
errJSON(w, 404, "not found")
|
|
}
|
|
}
|
|
// login gates on failure counts BEFORE verifying the password. While an IP is
|
|
// locked — or the global failure budget is spent — every attempt, including a
|
|
// correct password, gets the same 429 without a comparison, so responses never
|
|
// reveal whether a candidate password is right and brute force stays capped
|
|
// at loginFailuresPerMinute per IP and loginGlobalFailuresPerMinute overall.
|
|
// Successful logins are not failures: their gate slot is released by clearing
|
|
// the IP bucket, and only wrong passwords feed the global window.
|
|
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
|
|
ip := s.clientIP(r)
|
|
now := time.Now()
|
|
var x struct {
|
|
Password string `json:"password"`
|
|
}
|
|
if decode(w, r, &x) != nil {
|
|
errJSON(w, 400, "invalid request")
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
pruneFailures(s.attempts, now)
|
|
s.loginFailures = trimRecent(s.loginFailures, now)
|
|
recent := []time.Time{}
|
|
for _, t := range s.attempts[ip] {
|
|
if now.Sub(t) < time.Minute {
|
|
recent = append(recent, t)
|
|
}
|
|
}
|
|
blocked := len(recent) >= loginFailuresPerMinute || len(s.loginFailures) >= loginGlobalFailuresPerMinute || (len(recent) == 0 && len(s.attempts) >= 10000)
|
|
if !blocked {
|
|
// Reserve the gate slot atomically with the check so a concurrent
|
|
// burst cannot slip past the limit; a success erases the bucket below.
|
|
s.attempts[ip] = append(recent, now)
|
|
}
|
|
s.mu.Unlock()
|
|
if blocked {
|
|
w.Header().Set("Retry-After", "60")
|
|
errJSON(w, 429, "too many attempts")
|
|
return
|
|
}
|
|
h := sha256.Sum256([]byte(x.Password))
|
|
if subtle.ConstantTimeCompare(h[:], s.password[:]) != 1 {
|
|
s.mu.Lock()
|
|
s.loginFailures = append(s.loginFailures, now)
|
|
s.mu.Unlock()
|
|
errJSON(w, 401, "invalid credentials")
|
|
return
|
|
}
|
|
v := randomID()
|
|
s.mu.Lock()
|
|
delete(s.attempts, ip)
|
|
for k, t := range s.sessions {
|
|
if now.After(t) {
|
|
delete(s.sessions, k)
|
|
}
|
|
}
|
|
if c, e := r.Cookie("ai_router_session"); e == nil {
|
|
delete(s.sessions, c.Value)
|
|
}
|
|
s.sessions[v] = now.Add(24 * time.Hour)
|
|
s.mu.Unlock()
|
|
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", Value: v, Path: "/", MaxAge: 86400, HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
|
writeJSON(w, map[string]bool{"authenticated": true})
|
|
}
|