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 != "" {
			writeError(w, http.StatusNotFound, "not_found", "use GET /api/keys")
			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"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if body.Role == "" {
			body.Role = "user"
		}
		if body.Role != "admin" && body.Role != "user" {
			writeError(w, http.StatusBadRequest, "invalid_request", "role must be admin or user")
			return
		}
		rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
		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"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		rec, err := g.core.UpdateKey(path, body.Name, body.Role, body.Models, body.Note)
		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", "")
	}
}

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