Luminary/internal/handler/admin/ingress.go

package admin

import (
	"crypto/rand"
	"encoding/hex"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/rose_cat707/luminary/internal/auth"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/store/cache"
)

func (h *Handler) ListIngressKeys(c *gin.Context) {
	keys := h.cache.ListIngressKeys()
	c.JSON(http.StatusOK, gin.H{"data": keys})
}

func (h *Handler) CreateIngressKey(c *gin.Context) {
	var req struct {
		Name              string                  `json:"name"`
		Protocol          model.IngressKeyProtocol `json:"protocol"`
		BudgetLimit       float64                 `json:"budget_limit"`
		RateLimitRPM      int                     `json:"rate_limit_rpm"`
		RateLimitTPM      int                     `json:"rate_limit_tpm"`
		AllowedProviders  []string                `json:"allowed_providers"`
		AllowedModels     []string                `json:"allowed_models"`
		ExpiresAt         *time.Time              `json:"expires_at"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || req.Name == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
		return
	}
	protocol := req.Protocol
	if protocol == "" {
		protocol = model.IngressProtocolText
	}
	if !protocol.Valid() {
		c.JSON(http.StatusBadRequest, gin.H{"error": "protocol must be text or image"})
		return
	}
	plainKey, err := generateIngressKey()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "key generation failed"})
		return
	}
	hash := auth.HashIngressKey(plainKey)
	providersJSON := `["*"]`
	modelsJSON := `["*"]`
	if req.AllowedProviders != nil {
		if len(req.AllowedProviders) == 0 {
			providersJSON = `["*"]`
		} else {
			providersJSON = cache.ModelsToJSON(req.AllowedProviders)
		}
	}
	if req.AllowedModels != nil {
		if len(req.AllowedModels) == 0 {
			modelsJSON = `["*"]`
		} else {
			modelsJSON = cache.ModelsToJSON(req.AllowedModels)
		}
	}
	k := model.IngressKey{
		Name:                 req.Name,
		KeyHash:              hash,
		Prefix:               plainKey[:12] + "...",
		Protocol:             protocol,
		Enabled:              true,
		BudgetLimit:          req.BudgetLimit,
		RateLimitRPM:         req.RateLimitRPM,
		RateLimitTPM:         req.RateLimitTPM,
		AllowedProvidersJSON: providersJSON,
		AllowedModelsJSON:    modelsJSON,
		ExpiresAt:            req.ExpiresAt,
	}
	if err := h.db.DB().Create(&k).Error; err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	h.cache.SetIngressKey(k)
	c.JSON(http.StatusCreated, gin.H{"key": k, "plain_key": plainKey})
}

func (h *Handler) GetIngressKey(c *gin.Context) {
	id := parseUint(c.Param("id"))
	k, ok := h.cache.GetIngressKey(id)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.JSON(http.StatusOK, k)
}

func (h *Handler) UpdateIngressKey(c *gin.Context) {
	id := parseUint(c.Param("id"))
	k, ok := h.cache.GetIngressKey(id)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	var req struct {
		Name             string                   `json:"name"`
		Protocol         *model.IngressKeyProtocol `json:"protocol"`
		Enabled          *bool                    `json:"enabled"`
		BudgetLimit      *float64                 `json:"budget_limit"`
		RateLimitRPM     *int                     `json:"rate_limit_rpm"`
		RateLimitTPM     *int                     `json:"rate_limit_tpm"`
		AllowedProviders []string                 `json:"allowed_providers"`
		AllowedModels    []string                 `json:"allowed_models"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
		return
	}
	if req.Protocol != nil {
		if !req.Protocol.Valid() {
			c.JSON(http.StatusBadRequest, gin.H{"error": "protocol must be text or image"})
			return
		}
		k.Protocol = *req.Protocol
	}
	if req.Name != "" {
		k.Name = req.Name
	}
	if req.Enabled != nil {
		k.Enabled = *req.Enabled
	}
	if req.BudgetLimit != nil {
		k.BudgetLimit = *req.BudgetLimit
	}
	if req.RateLimitRPM != nil {
		k.RateLimitRPM = *req.RateLimitRPM
	}
	if req.RateLimitTPM != nil {
		k.RateLimitTPM = *req.RateLimitTPM
	}
	if req.AllowedProviders != nil {
		if len(req.AllowedProviders) == 0 {
			k.AllowedProvidersJSON = `["*"]`
		} else {
			k.AllowedProvidersJSON = cache.ModelsToJSON(req.AllowedProviders)
		}
	}
	if req.AllowedModels != nil {
		if len(req.AllowedModels) == 0 {
			k.AllowedModelsJSON = `["*"]`
		} else {
			k.AllowedModelsJSON = cache.ModelsToJSON(req.AllowedModels)
		}
	}
	h.db.DB().Save(&k)
	h.cache.SetIngressKey(k)
	c.JSON(http.StatusOK, k)
}

func (h *Handler) DeleteIngressKey(c *gin.Context) {
	id := parseUint(c.Param("id"))
	h.db.DB().Where("ingress_key_id = ?", id).Delete(&model.RoutingRule{})
	h.db.DB().Delete(&model.IngressKey{}, id)
	h.cache.DeleteIngressKey(id)
	h.cache.DeleteRoutingRulesForIngress(id)
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (h *Handler) IngressUsage(c *gin.Context) {
	id := parseUint(c.Param("id"))
	var records []model.UsageRecord
	h.db.DB().Where("ingress_key_id = ?", id).Order("created_at desc").Limit(50).Find(&records)
	k, _ := h.cache.GetIngressKey(id)
	c.JSON(http.StatusOK, gin.H{"key": k, "recent": records})
}

func generateIngressKey() (string, error) {
	b := make([]byte, 24)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return "sk-lum-" + hex.EncodeToString(b), nil
}