ModelRouter/internal/gateway/keys.go

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", "")
	}
}