From 5a9a6d54a399263318d4ad5ea24b8c4c41af8fb2 Mon Sep 17 00:00:00 2001 From: Ray Andrew Date: Sun, 15 Feb 2026 02:27:43 -0600 Subject: [PATCH] feat(gateway): fix static token initialized in DB --- llm-gateway/cmd/gateway/main.go | 45 ++++------ llm-gateway/internal/auth/store.go | 88 +++++++++++-------- .../dashboard/templates/partials/tokens.html | 27 ++++-- llm-gateway/internal/proxy/auth.go | 6 +- 4 files changed, 94 insertions(+), 72 deletions(-) diff --git a/llm-gateway/cmd/gateway/main.go b/llm-gateway/cmd/gateway/main.go index cb33f18..a418431 100644 --- a/llm-gateway/cmd/gateway/main.go +++ b/llm-gateway/cmd/gateway/main.go @@ -90,13 +90,25 @@ func main() { } log.Printf("Registered %d models", len(cfg.Models)) - // Auth store - authStore := auth.NewStore(db.DB) + // Auth store (static tokens checked in-memory, not seeded to DB) + var staticTokens []auth.StaticToken + for _, t := range cfg.Tokens { + if t.Key != "" { + staticTokens = append(staticTokens, auth.StaticToken{ + Name: t.Name, + Key: t.Key, + RateLimitRPM: t.RateLimitRPM, + DailyBudgetUSD: t.DailyBudgetUSD, + }) + log.Printf("Loaded static token: %s", t.Name) + } + } + authStore := auth.NewStore(db.DB, staticTokens) authMiddleware := auth.NewMiddleware(authStore) authHandlers := auth.NewHandlers(authStore, cfg.Server.SessionSecret) - // Seed default admin and static tokens - seedAdminAndTokens(cfg, authStore) + // Seed default admin + seedDefaultAdmin(cfg, authStore) // Metrics m := metrics.New() @@ -244,9 +256,8 @@ func main() { log.Println("Stopped") } -// seedAdminAndTokens creates the default admin and seeds static tokens from config. -func seedAdminAndTokens(cfg *config.Config, authStore *auth.Store) { - // Seed default admin if no users exist +// seedDefaultAdmin creates the default admin user if no users exist. +func seedDefaultAdmin(cfg *config.Config, authStore *auth.Store) { if !authStore.HasAnyUser() { da := cfg.Server.DefaultAdmin if da.Username != "" && da.Password != "" { @@ -258,24 +269,4 @@ func seedAdminAndTokens(cfg *config.Config, authStore *auth.Store) { } } } - - // Seed static tokens from config - if len(cfg.Tokens) > 0 { - admin, err := authStore.GetFirstAdmin() - if err != nil { - log.Printf("WARNING: no admin user found, cannot seed static tokens") - return - } - - for _, t := range cfg.Tokens { - if t.Key == "" { - continue - } - if err := authStore.SeedStaticToken(admin.ID, t.Name, t.Key, t.RateLimitRPM, t.DailyBudgetUSD); err != nil { - log.Printf("WARNING: failed to seed token %q: %v", t.Name, err) - } else { - log.Printf("Seeded static token: %s", t.Name) - } - } - } } diff --git a/llm-gateway/internal/auth/store.go b/llm-gateway/internal/auth/store.go index c829444..170af0c 100644 --- a/llm-gateway/internal/auth/store.go +++ b/llm-gateway/internal/auth/store.go @@ -42,12 +42,21 @@ type APIToken struct { LastUsedAt int64 `json:"last_used_at"` } -type Store struct { - db *sql.DB +// StaticToken represents a token defined in config (checked in-memory, never stored in DB). +type StaticToken struct { + Name string + Key string + RateLimitRPM int + DailyBudgetUSD float64 } -func NewStore(db *sql.DB) *Store { - return &Store{db: db} +type Store struct { + db *sql.DB + staticTokens []StaticToken +} + +func NewStore(db *sql.DB, staticTokens []StaticToken) *Store { + return &Store{db: db, staticTokens: staticTokens} } func (s *Store) HasAnyUser() bool { @@ -100,12 +109,6 @@ func (s *Store) GetUserByID(id int64) (*User, error) { )) } -func (s *Store) GetFirstAdmin() (*User, error) { - return s.scanUser(s.db.QueryRow( - "SELECT id, username, email, password_hash, is_admin, totp_secret, totp_enabled, created_at, updated_at FROM users WHERE is_admin = 1 ORDER BY id LIMIT 1", - )) -} - func (s *Store) scanUser(row *sql.Row) (*User, error) { var u User var isAdmin, totpEnabled int @@ -271,6 +274,24 @@ func (s *Store) CreateAPIToken(userID int64, name string, rateLimitRPM int, dail } func (s *Store) LookupAPIToken(key string) (*APIToken, error) { + // Check static tokens first (from config, never stored in DB) + for _, st := range s.staticTokens { + if st.Key == key { + prefix := st.Key + if len(prefix) > 11 { + prefix = prefix[:11] + } + return &APIToken{ + ID: -1, // sentinel: static token + Name: st.Name, + KeyPrefix: prefix, + RateLimitRPM: st.RateLimitRPM, + DailyBudgetUSD: st.DailyBudgetUSD, + }, nil + } + } + + // Fall back to DB tokens hash := sha256.Sum256([]byte(key)) keyHash := hex.EncodeToString(hash[:]) @@ -286,6 +307,23 @@ func (s *Store) LookupAPIToken(key string) (*APIToken, error) { } func (s *Store) ListAPITokens(userID int64) ([]APIToken, error) { + // Include static tokens (shown for all users, not deletable) + var tokens []APIToken + for _, st := range s.staticTokens { + prefix := st.Key + if len(prefix) > 11 { + prefix = prefix[:11] + } + tokens = append(tokens, APIToken{ + ID: -1, // sentinel: static token + Name: st.Name, + KeyPrefix: prefix, + RateLimitRPM: st.RateLimitRPM, + DailyBudgetUSD: st.DailyBudgetUSD, + }) + } + + // DB tokens var rows *sql.Rows var err error if userID == 0 { @@ -295,15 +333,14 @@ func (s *Store) ListAPITokens(userID int64) ([]APIToken, error) { rows, err = s.db.Query("SELECT id, name, key_hash, key_prefix, user_id, rate_limit_rpm, daily_budget_usd, created_at, last_used_at FROM api_tokens WHERE user_id = ? ORDER BY id", userID) } if err != nil { - return nil, err + return tokens, nil // return static tokens even if DB query fails } defer rows.Close() - var tokens []APIToken for rows.Next() { var t APIToken if err := rows.Scan(&t.ID, &t.Name, &t.KeyHash, &t.KeyPrefix, &t.UserID, &t.RateLimitRPM, &t.DailyBudgetUSD, &t.CreatedAt, &t.LastUsedAt); err != nil { - return nil, err + return tokens, nil } tokens = append(tokens, t) } @@ -331,31 +368,6 @@ func (s *Store) UpdateAPITokenLastUsed(id int64) { s.db.Exec("UPDATE api_tokens SET last_used_at = ? WHERE id = ?", time.Now().Unix(), id) } -// SeedStaticToken creates a token by name if it doesn't already exist (idempotent). -func (s *Store) SeedStaticToken(userID int64, name, plainKey string, rateLimitRPM int, dailyBudgetUSD float64) error { - // Check if token with this name already exists - var count int - s.db.QueryRow("SELECT COUNT(*) FROM api_tokens WHERE name = ?", name).Scan(&count) - if count > 0 { - return nil // already seeded - } - - keyPrefix := plainKey - if len(keyPrefix) > 11 { - keyPrefix = keyPrefix[:11] - } - - hash := sha256.Sum256([]byte(plainKey)) - keyHash := hex.EncodeToString(hash[:]) - - now := time.Now().Unix() - _, err := s.db.Exec( - "INSERT INTO api_tokens (name, key_hash, key_prefix, user_id, rate_limit_rpm, daily_budget_usd, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", - name, keyHash, keyPrefix, userID, rateLimitRPM, dailyBudgetUSD, now, - ) - return err -} - func (s *Store) UpdateUsername(userID int64, newUsername string) error { _, err := s.db.Exec("UPDATE users SET username = ?, updated_at = ? WHERE id = ?", newUsername, time.Now().Unix(), userID) return err diff --git a/llm-gateway/internal/dashboard/templates/partials/tokens.html b/llm-gateway/internal/dashboard/templates/partials/tokens.html index 550616a..58ebbbf 100644 --- a/llm-gateway/internal/dashboard/templates/partials/tokens.html +++ b/llm-gateway/internal/dashboard/templates/partials/tokens.html @@ -10,10 +10,30 @@
+

Static Tokens (from config, managed via environment variables)

+ + + + {{range .Tokens}}{{if lt .ID 0}} + + + + + + + + {{end}}{{end}} + +
NamePrefixRate LimitBudget
{{.Name}}{{.KeyPrefix}}...{{if eq .RateLimitRPM 0}}unlimited{{else}}{{.RateLimitRPM}} rpm{{end}}{{if gt .DailyBudgetUSD 0.0}}${{printf "%.2f" .DailyBudgetUSD}}{{else}}unlimited{{end}}config
+
+ +
+

Dynamic Tokens (created via dashboard)

- {{range .Tokens}} + {{$hasDynamic := false}} + {{range .Tokens}}{{if gt .ID 0}} @@ -23,9 +43,7 @@ - {{else}} - - {{end}} + {{end}}{{end}}
NamePrefixRate LimitBudgetCreatedLast Used
{{.Name}} {{.KeyPrefix}}...{{if gt .LastUsedAt 0}}{{formatTime .LastUsedAt}}{{else}}never{{end}}
No API tokens yet. Create one to get started.
@@ -88,7 +106,6 @@ async function doCreateToken(e) { closeModal(); document.getElementById('new-token-key').textContent = data.key; document.getElementById('new-token-display').style.display = 'block'; - // Reload tokens partial htmx.ajax('GET', '/tokens', {target: '#content', swap: 'innerHTML'}); } catch (e) { document.getElementById('create-token-error').innerHTML = '
' + e.message + '
'; } } diff --git a/llm-gateway/internal/proxy/auth.go b/llm-gateway/internal/proxy/auth.go index b885550..ca0944f 100644 --- a/llm-gateway/internal/proxy/auth.go +++ b/llm-gateway/internal/proxy/auth.go @@ -31,8 +31,10 @@ func (a *AuthMiddleware) Authenticate(next http.Handler) http.Handler { return } - // Update last used asynchronously - go a.authStore.UpdateAPITokenLastUsed(token.ID) + // Update last used asynchronously (skip for static tokens) + if token.ID > 0 { + go a.authStore.UpdateAPITokenLastUsed(token.ID) + } ctx := withTokenName(r.Context(), token.Name) ctx = withAPIToken(ctx, token)