ai-gateway-go/internal/apikey/admin_http.go

package apikey

import (
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

	"aigateway.local/core/internal/identity"
	"aigateway.local/core/internal/platform/apiresponse"
)

type AdminHTTPHandler struct {
	repository    *Repository
	authenticator *Authenticator
	usage         *UsageStore
	identity      *identity.Service
	mux           *http.ServeMux
}

type createRequest struct {
	Name                string     `json:"name"`
	Scopes              []string   `json:"scopes"`
	RequestsPerMinute   int        `json:"requests_per_minute"`
	MonthlyRequestQuota int64      `json:"monthly_request_quota"`
	MonthlyTokenQuota   int64      `json:"monthly_token_quota"`
	ExpiresAt           *time.Time `json:"expires_at"`
}

type limitsRequest struct {
	RequestsPerMinute   int   `json:"requests_per_minute"`
	MonthlyRequestQuota int64 `json:"monthly_request_quota"`
	MonthlyTokenQuota   int64 `json:"monthly_token_quota"`
}

func NewAdminHTTPHandler(repository *Repository, authenticator *Authenticator, identityService *identity.Service) *AdminHTTPHandler {
	h := &AdminHTTPHandler{repository: repository, authenticator: authenticator, identity: identityService, mux: http.NewServeMux()}
	h.mux.HandleFunc("GET /api/v1/admin/api-keys", h.list)
	h.mux.HandleFunc("POST /api/v1/admin/api-keys", h.create)
	h.mux.HandleFunc("PUT /api/v1/admin/api-keys/{api_key_id}/limits", h.updateLimits)
	h.mux.HandleFunc("DELETE /api/v1/admin/api-keys/{api_key_id}", h.revoke)
	return h
}

func (h *AdminHTTPHandler) SetUsageStore(store *UsageStore) { h.usage = store }

func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	h.mux.ServeHTTP(writer, request)
}

func (h *AdminHTTPHandler) list(writer http.ResponseWriter, request *http.Request) {
	if _, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyRead); !ok {
		return
	}
	records, err := h.repository.List(request.Context())
	if err != nil {
		h.writeError(writer, err)
		return
	}
	ids := make([]string, 0, len(records))
	for _, record := range records {
		ids = append(ids, record.ID)
	}
	usage, err := h.usage.MonthlyTokens(request.Context(), ids, time.Now())
	if err != nil {
		h.writeError(writer, err)
		return
	}
	items := make([]map[string]any, 0, len(records))
	for _, record := range records {
		item := publicRecord(record)
		item["monthly_token_usage"] = usage[record.ID]
		items = append(items, item)
	}
	apiresponse.OK(writer, items)
}

func (h *AdminHTTPHandler) updateLimits(writer http.ResponseWriter, request *http.Request) {
	account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
	if !ok {
		return
	}
	var input limitsRequest
	decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&input) != nil || !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
		apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
		return
	}
	record, hash, err := h.repository.UpdateLimits(request.Context(), request.PathValue("api_key_id"), input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, account.ID)
	if err != nil {
		h.writeError(writer, err)
		return
	}
	if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
		h.writeError(writer, err)
		return
	}
	response := publicRecord(record)
	if h.usage != nil {
		usage, usageErr := h.usage.MonthlyTokens(request.Context(), []string{record.ID}, time.Now())
		if usageErr == nil {
			response["monthly_token_usage"] = usage[record.ID]
		}
	}
	apiresponse.OK(writer, response)
}

func validLimits(requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64) bool {
	return requestsPerMinute >= 0 && requestsPerMinute <= 1_000_000 &&
		monthlyRequestQuota >= 0 && monthlyRequestQuota <= 1_000_000_000_000 &&
		monthlyTokenQuota >= 0 && monthlyTokenQuota <= 1_000_000_000_000_000
}

func (h *AdminHTTPHandler) create(writer http.ResponseWriter, request *http.Request) {
	account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
	if !ok {
		return
	}
	var input createRequest
	decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&input) != nil {
		apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
		return
	}
	input.Name = strings.TrimSpace(input.Name)
	if input.Name == "" || len(input.Name) > 128 || len(input.Scopes) == 0 || input.ExpiresAt != nil && !input.ExpiresAt.After(time.Now()) {
		apiresponse.Error(writer, http.StatusBadRequest, "名称、权限范围或过期时间无效")
		return
	}
	if !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
		apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
		return
	}
	for _, scope := range input.Scopes {
		if scope != "gateway:invoke" && scope != "*" {
			apiresponse.Error(writer, http.StatusBadRequest, "包含不支持的权限范围")
			return
		}
	}
	record, secret, err := h.repository.Create(request.Context(), input.Name, input.Scopes, input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, input.ExpiresAt, account.ID)
	if err != nil {
		h.writeError(writer, err)
		return
	}
	response := publicRecord(record)
	response["key"] = secret
	response["warning"] = "密钥只显示一次，请立即安全保存"
	apiresponse.OK(writer, response)
}

func (h *AdminHTTPHandler) revoke(writer http.ResponseWriter, request *http.Request) {
	account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
	if !ok {
		return
	}
	hash, err := h.repository.Revoke(request.Context(), request.PathValue("api_key_id"), account.ID)
	if err != nil {
		h.writeError(writer, err)
		return
	}
	if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
		h.writeError(writer, err)
		return
	}
	apiresponse.OK(writer, map[string]bool{"revoked": true})
}

func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) (identity.Account, bool) {
	account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
	if err != nil {
		apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
		return identity.Account{}, false
	}
	if !identity.HasPermission(account, permission) {
		apiresponse.Error(writer, http.StatusForbidden, "缺少 API Key 操作权限")
		return identity.Account{}, false
	}
	return account, true
}

func (h *AdminHTTPHandler) writeError(writer http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, ErrInvalid):
		apiresponse.Error(writer, http.StatusNotFound, "API Key 不存在或已撤销")
	case errors.Is(err, ErrStore):
		apiresponse.Error(writer, http.StatusServiceUnavailable, "API Key 服务暂不可用")
	default:
		apiresponse.Error(writer, http.StatusInternalServerError, "API Key 处理失败")
	}
}

func publicRecord(record Record) map[string]any {
	return map[string]any{
		"id": record.ID, "name": record.Name, "key_prefix": record.KeyPrefix,
		"scopes": record.Scopes, "enabled": record.Enabled, "expires_at": record.ExpiresAt,
		"requests_per_minute": record.RequestsPerMinute, "monthly_request_quota": record.MonthlyRequestQuota,
		"monthly_token_quota": record.MonthlyTokenQuota,
		"last_used_at":        record.LastUsedAt, "created_at": record.CreatedAt,
	}
}