package gateway import ( "context" "encoding/json" "net/http" "strings" "llmsproxy/internal/config" ) // handleKeysAPI manages gateway keys: GET /api/keys (admin: all keys), // GET /api/keys/me (own key for any role), POST /api/keys (admin: create), // PUT /api/keys/{key} (admin: update), DELETE /api/keys/{key} (admin: remove). func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) { path := strings.TrimPrefix(r.URL.Path, "/api/keys") path = strings.Trim(path, "/") role := reqRole(r.Context()) if path == "me" { g.handleKeyMe(w, r) return } if role != "admin" { writeError(w, http.StatusForbidden, "forbidden", "admin role required") return } switch r.Method { case http.MethodGet: if path != "" { // A GET on a specific key is almost always a client that meant to // DELETE or PUT it but let fetch default to GET. Name the verbs // instead of only pointing back at the collection endpoint. writeError(w, http.StatusNotFound, "not_found", "no such endpoint; use GET /api/keys to list, DELETE /api/keys/{key} to remove, PUT /api/keys/{key} to update") return } writeJSON(w, http.StatusOK, map[string]interface{}{"keys": g.core.ListKeys()}) case http.MethodPost: var body struct { Name string `json:"name"` Role string `json:"role"` Models []config.ModelScope `json:"models"` Note string `json:"note"` TokenQuota *int64 `json:"token_quota"` ReqQuota *int64 `json:"req_quota"` Period *string `json:"period"` Hours *int64 `json:"hours"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error()) return } body.Role = config.NormalizeRole(body.Role) q := config.KeyQuota{ TokenQuota: optInt64(body.TokenQuota), ReqQuota: optInt64(body.ReqQuota), Period: optString(body.Period), Hours: optInt64(body.Hours), } rec, err := g.core.CreateKeyWithQuota(body.Name, body.Role, body.Models, body.Note, q) if err != nil { writeError(w, http.StatusBadRequest, "key_error", err.Error()) return } writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "key": rec}) case http.MethodPut, http.MethodPatch: if path == "" { writeError(w, http.StatusBadRequest, "invalid_request", "key required") return } var body struct { Name string `json:"name"` Role string `json:"role"` Models []config.ModelScope `json:"models"` Note string `json:"note"` TokenQuota *int64 `json:"token_quota"` ReqQuota *int64 `json:"req_quota"` Period *string `json:"period"` Hours *int64 `json:"hours"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error()) return } // Quota fields are pointers so "absent" is distinguishable from // "set to 0": omitting them leaves the stored caps alone, while // sending 0 explicitly lifts a cap. Without this, editing only the // model scope would silently clear a key's budget. var q *config.KeyQuota if body.TokenQuota != nil || body.ReqQuota != nil || body.Period != nil || body.Hours != nil { q = &config.KeyQuota{ TokenQuota: optInt64(body.TokenQuota), ReqQuota: optInt64(body.ReqQuota), Period: optString(body.Period), Hours: optInt64(body.Hours), } } rec, err := g.core.UpdateKeyWithQuota(path, body.Name, body.Role, body.Models, body.Note, q) if err != nil { writeError(w, http.StatusBadRequest, "key_error", err.Error()) return } writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "key": rec}) case http.MethodDelete: if path == "" { writeError(w, http.StatusBadRequest, "invalid_request", "key required") return } if path == reqKey(r.Context()) { writeError(w, http.StatusBadRequest, "invalid_request", "cannot delete the key you are logged in with") return } ok, err := g.core.DeleteKey(path) if err != nil { writeError(w, http.StatusBadRequest, "key_error", err.Error()) return } if !ok { writeError(w, http.StatusNotFound, "not_found", "key not found") return } writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true}) default: writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "") } } // optInt64 dereferences an optional quota field, treating absent as 0. func optInt64(p *int64) int64 { if p == nil { return 0 } return *p } // optString dereferences an optional quota field, treating absent as "". func optString(p *string) string { if p == nil { return "" } return *p } // handleKeyMe returns the authenticated key's own record (users see only // themselves; admins can use this as a convenience too). func (g *Gateway) handleKeyMe(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET") return } rec, ok := g.core.FindKey(reqKey(r.Context())) if !ok { writeError(w, http.StatusUnauthorized, "invalid_api_key", "key not found") return } if !rec.Seed { for _, s := range g.core.GatewayKeys() { if s == rec.Key { rec.Seed = true break } } } writeJSON(w, http.StatusOK, map[string]interface{}{"key": rec}) } // allowedModels returns the model scope for the request's key; nil means // unrestricted (admin keys and user keys without an explicit scope). func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope { if reqRole(ctx) == "admin" { return nil } rec, ok := g.core.FindKey(reqKey(ctx)) if !ok { return nil } return rec.Models } // handleAutoAPI manages the AUTO scheduling slots: GET /api/auto returns the // current rules; PUT /api/auto replaces them (admin only). func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodGet { writeJSON(w, http.StatusOK, map[string]interface{}{ "rules": g.core.AutoRules(), "image_rules": g.core.AutoImageRules(), "states": g.core.AutoSlotStates(), }) return } if reqRole(r.Context()) != "admin" { writeError(w, http.StatusForbidden, "forbidden", "admin role required") return } switch r.Method { case http.MethodPut, http.MethodPost: var body struct { Rules []config.ModelScope `json:"rules"` ImageRules []config.ModelScope `json:"image_rules"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error()) return } if body.Rules != nil { if err := g.core.SaveAutoRules(body.Rules); err != nil { writeError(w, http.StatusBadRequest, "auto_error", err.Error()) return } } if body.ImageRules != nil { if err := g.core.SaveAutoImageRules(body.ImageRules); err != nil { writeError(w, http.StatusBadRequest, "auto_error", err.Error()) return } } writeJSON(w, http.StatusOK, map[string]interface{}{ "ok": true, "rules": g.core.AutoRules(), "image_rules": g.core.AutoImageRules(), }) default: writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "") } }