rayandrew/ai-servers
rayandrew/ai-servers
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(gateway): fix static token initialized in DB

Browse source
This commit is contained in:
Ray Andrew 2026-02-15 02:27:43 -06:00
parent 9dea812a47
commit 5a9a6d54a3
Signed by: rayandrew
SSH key fingerprint: SHA256:EUCV+qCSqkap8rR+p+zGjxHfKI06G0GJKgo1DIOniQY
4 changed files with 94 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
llm-gateway/cmd/gateway/main.go
View file

@ -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)
}
}
}
}

88
llm-gateway/internal/auth/store.go
View file

@ -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

27
llm-gateway/internal/dashboard/templates/partials/tokens.html
View file

@ -10,10 +10,30 @@
</div>
<div class="section">
<h2>Static Tokens <span style="font-size:0.75rem;color:#64748b;font-weight:400;">(from config, managed via environment variables)</span></h2>
<table>
<thead><tr><th>Name</th><th>Prefix</th><th>Rate Limit</th><th>Budget</th><th></th></tr></thead>
<tbody>
{{range .Tokens}}{{if lt .ID 0}}
<tr>
<td>{{.Name}}</td>
<td><code>{{.KeyPrefix}}...</code></td>
<td>{{if eq .RateLimitRPM 0}}unlimited{{else}}{{.RateLimitRPM}} rpm{{end}}</td>
<td>{{if gt .DailyBudgetUSD 0.0}}${{printf "%.2f" .DailyBudgetUSD}}{{else}}unlimited{{end}}</td>
<td><span class="badge badge-totp">config</span></td>
</tr>
{{end}}{{end}}
</tbody>
</table>
</div>
<div class="section">
<h2>Dynamic Tokens <span style="font-size:0.75rem;color:#64748b;font-weight:400;">(created via dashboard)</span></h2>
<table>
<thead><tr><th>Name</th><th>Prefix</th><th>Rate Limit</th><th>Budget</th><th>Created</th><th>Last Used</th><th></th></tr></thead>
<tbody id="tokens-tbody">
{{range .Tokens}}
{{$hasDynamic := false}}
{{range .Tokens}}{{if gt .ID 0}}
<tr>
<td>{{.Name}}</td>
<td><code>{{.KeyPrefix}}...</code></td>
@ -23,9 +43,7 @@
<td>{{if gt .LastUsedAt 0}}{{formatTime .LastUsedAt}}{{else}}never{{end}}</td>
<td><button class="btn btn-sm btn-danger" onclick="deleteToken({{.ID}})">Revoke</button></td>
</tr>
{{else}}
<tr><td colspan="7" style="color:#64748b;text-align:center;padding:20px;">No API tokens yet. Create one to get started.</td></tr>
{{end}}
{{end}}{{end}}
</tbody>
</table>
</div>
@ -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 = '<div class="error-msg">' + e.message + '</div>'; }
}

4
llm-gateway/internal/proxy/auth.go
View file

@ -31,8 +31,10 @@ func (a *AuthMiddleware) Authenticate(next http.Handler) http.Handler {
return
}
// Update last used asynchronously
// 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)