package cache import ( "encoding/json" "sync" "sync/atomic" "time" "github.com/rose_cat707/luminary/internal/auth" "github.com/rose_cat707/luminary/internal/model" ) type IngressKeySnapshot struct { Key model.IngressKey KeyPlain string // only set on create, not persisted in cache long-term } type UsageDelta struct { IngressKeyID uint ProviderKeyID uint ProviderID uint ClientIP string Endpoint string Model string TokensIn int TokensOut int Cost float64 LatencyMs int64 Status string Stream bool ErrorMsg string RequestBody string ResponseBody string BudgetUsed float64 Requests int64 Tokens int64 } type Store struct { mu sync.RWMutex providers map[uint]*model.Provider providerKeys map[uint]*model.ProviderKey providerByName map[string]*model.Provider models map[uint][]model.ModelEntry ingressKeys map[uint]*model.IngressKey ingressByHash map[string]*model.IngressKey ipRules []model.IPRule routingRules map[uint][]model.RoutingRule // ingressKeyID -> rules healthStatus map[uint]model.HealthCheckRecord usageDeltas []UsageDelta usageMu sync.Mutex // rate limit counters: ingressKeyID -> window start + counts rateMu sync.Mutex rateWindows map[uint]*rateWindow } type rateWindow struct { start time.Time requests int tokens int } func New() *Store { return &Store{ providers: make(map[uint]*model.Provider), providerKeys: make(map[uint]*model.ProviderKey), providerByName: make(map[string]*model.Provider), models: make(map[uint][]model.ModelEntry), ingressKeys: make(map[uint]*model.IngressKey), ingressByHash: make(map[string]*model.IngressKey), routingRules: make(map[uint][]model.RoutingRule), healthStatus: make(map[uint]model.HealthCheckRecord), rateWindows: make(map[uint]*rateWindow), } } func (c *Store) LoadProviders(providers []model.Provider) { c.mu.Lock() defer c.mu.Unlock() c.providers = make(map[uint]*model.Provider) c.providerByName = make(map[string]*model.Provider) for i := range providers { p := providers[i] c.providers[p.ID] = &p c.providerByName[p.Name] = &p } } func (c *Store) LoadProviderKeys(keys []model.ProviderKey) { c.mu.Lock() defer c.mu.Unlock() c.providerKeys = make(map[uint]*model.ProviderKey) for i := range keys { k := keys[i] c.providerKeys[k.ID] = &k } } func (c *Store) LoadModels(models []model.ModelEntry) { c.mu.Lock() defer c.mu.Unlock() c.models = make(map[uint][]model.ModelEntry) for _, m := range models { c.models[m.ProviderID] = append(c.models[m.ProviderID], m) } } func (c *Store) LoadIngressKeys(keys []model.IngressKey) { c.mu.Lock() defer c.mu.Unlock() c.ingressKeys = make(map[uint]*model.IngressKey) c.ingressByHash = make(map[string]*model.IngressKey) for i := range keys { k := keys[i] c.ingressKeys[k.ID] = &k c.ingressByHash[k.KeyHash] = &k } } func (c *Store) LoadIPRules(rules []model.IPRule) { c.mu.Lock() defer c.mu.Unlock() c.ipRules = rules } func (c *Store) GetProvider(id uint) (model.Provider, bool) { c.mu.RLock() defer c.mu.RUnlock() p, ok := c.providers[id] if !ok { return model.Provider{}, false } return *p, true } func (c *Store) GetProviderByName(name string) (model.Provider, bool) { c.mu.RLock() defer c.mu.RUnlock() p, ok := c.providerByName[name] if !ok { return model.Provider{}, false } return *p, true } func (c *Store) ListProviders() []model.Provider { c.mu.RLock() defer c.mu.RUnlock() out := make([]model.Provider, 0, len(c.providers)) for _, p := range c.providers { out = append(out, *p) } return out } func (c *Store) SetProvider(p model.Provider) { c.mu.Lock() defer c.mu.Unlock() c.providers[p.ID] = &p c.providerByName[p.Name] = &p } func (c *Store) DeleteProvider(id uint) { c.mu.Lock() defer c.mu.Unlock() if p, ok := c.providers[id]; ok { delete(c.providerByName, p.Name) } delete(c.providers, id) delete(c.models, id) } func (c *Store) GetProviderKey(id uint) (model.ProviderKey, bool) { c.mu.RLock() defer c.mu.RUnlock() k, ok := c.providerKeys[id] if !ok { return model.ProviderKey{}, false } return *k, true } func (c *Store) ListProviderKeys(providerID uint) []model.ProviderKey { c.mu.RLock() defer c.mu.RUnlock() var out []model.ProviderKey for _, k := range c.providerKeys { if k.ProviderID == providerID { out = append(out, *k) } } return out } func (c *Store) ListAllProviderKeys() []model.ProviderKey { c.mu.RLock() defer c.mu.RUnlock() out := make([]model.ProviderKey, 0, len(c.providerKeys)) for _, k := range c.providerKeys { out = append(out, *k) } return out } func (c *Store) SetProviderKey(k model.ProviderKey) { c.mu.Lock() defer c.mu.Unlock() c.providerKeys[k.ID] = &k } func (c *Store) DeleteProviderKey(id uint) { c.mu.Lock() defer c.mu.Unlock() delete(c.providerKeys, id) } func (c *Store) GetModels(providerID uint) []model.ModelEntry { c.mu.RLock() defer c.mu.RUnlock() models := c.models[providerID] if models == nil { return []model.ModelEntry{} } return models } func (c *Store) SetModels(providerID uint, models []model.ModelEntry) { c.mu.Lock() defer c.mu.Unlock() c.models[providerID] = models } func (c *Store) ListIngressKeys() []model.IngressKey { c.mu.RLock() defer c.mu.RUnlock() out := make([]model.IngressKey, 0, len(c.ingressKeys)) for _, k := range c.ingressKeys { out = append(out, *k) } return out } func (c *Store) GetIngressKey(id uint) (model.IngressKey, bool) { c.mu.RLock() defer c.mu.RUnlock() k, ok := c.ingressKeys[id] if !ok { return model.IngressKey{}, false } return *k, true } func (c *Store) SetIngressKey(k model.IngressKey) { c.mu.Lock() defer c.mu.Unlock() c.ingressKeys[k.ID] = &k c.ingressByHash[k.KeyHash] = &k } func (c *Store) DeleteIngressKey(id uint) { c.mu.Lock() defer c.mu.Unlock() if k, ok := c.ingressKeys[id]; ok { delete(c.ingressByHash, k.KeyHash) } delete(c.ingressKeys, id) } func (c *Store) LookupIngressKey(plain string) (*model.IngressKey, bool) { c.mu.RLock() defer c.mu.RUnlock() hash := auth.HashIngressKey(plain) if k, ok := c.ingressByHash[hash]; ok && k.Enabled { cp := *k return &cp, true } for _, k := range c.ingressKeys { if !k.Enabled || !auth.IsLegacyIngressHash(k.KeyHash) { continue } if auth.CheckIngressKey(k.KeyHash, plain) { cp := *k return &cp, true } } return nil, false } func (c *Store) ResetProviderKeyDailyCounts() { c.mu.Lock() defer c.mu.Unlock() for _, k := range c.providerKeys { k.RequestsToday = 0 k.TokensToday = 0 } } func (c *Store) ListIPRules() []model.IPRule { c.mu.RLock() defer c.mu.RUnlock() out := make([]model.IPRule, len(c.ipRules)) copy(out, c.ipRules) return out } func (c *Store) SetIPRules(rules []model.IPRule) { c.mu.Lock() defer c.mu.Unlock() c.ipRules = rules } func (c *Store) SetHealthStatus(providerID uint, rec model.HealthCheckRecord) { c.mu.Lock() defer c.mu.Unlock() c.healthStatus[providerID] = rec if p, ok := c.providers[providerID]; ok { p.Status = rec.Status } } func (c *Store) GetHealthStatus() map[uint]model.HealthCheckRecord { c.mu.RLock() defer c.mu.RUnlock() out := make(map[uint]model.HealthCheckRecord, len(c.healthStatus)) for k, v := range c.healthStatus { out[k] = v } return out } func (c *Store) RecordUsage(delta UsageDelta) { c.mu.Lock() if k, ok := c.ingressKeys[delta.IngressKeyID]; ok { k.BudgetUsed += delta.BudgetUsed k.RequestCount += delta.Requests k.TokenCount += delta.Tokens } if pk, ok := c.providerKeys[delta.ProviderKeyID]; ok { pk.RequestsToday += int(delta.Requests) pk.TokensToday += delta.Tokens } c.mu.Unlock() c.usageMu.Lock() c.usageDeltas = append(c.usageDeltas, delta) c.usageMu.Unlock() } func (c *Store) DrainUsageDeltas() []UsageDelta { c.usageMu.Lock() defer c.usageMu.Unlock() if len(c.usageDeltas) == 0 { return nil } out := c.usageDeltas c.usageDeltas = nil return out } func (c *Store) CheckRateLimit(ingressKeyID uint, rpm, tpm int, tokens int) bool { if rpm <= 0 && tpm <= 0 { return true } c.rateMu.Lock() defer c.rateMu.Unlock() now := time.Now() w, ok := c.rateWindows[ingressKeyID] if !ok || now.Sub(w.start) >= time.Minute { c.rateWindows[ingressKeyID] = &rateWindow{start: now, requests: 1, tokens: tokens} return true } if rpm > 0 && w.requests >= rpm { return false } if tpm > 0 && w.tokens+tokens > tpm { return false } w.requests++ w.tokens += tokens return true } func ParseJSONList(s string) []string { var out []string if s == "" { return []string{"*"} } _ = json.Unmarshal([]byte(s), &out) if len(out) == 0 { return []string{"*"} } return out } func MatchesList(list []string, value string) bool { for _, item := range list { if item == "*" || item == value { return true } } return false } func ModelsToJSON(models []string) string { b, _ := json.Marshal(models) return string(b) } func (c *Store) LoadRoutingRules(rules []model.RoutingRule) { c.mu.Lock() defer c.mu.Unlock() c.routingRules = make(map[uint][]model.RoutingRule) for _, r := range rules { c.routingRules[r.IngressKeyID] = append(c.routingRules[r.IngressKeyID], r) } } func (c *Store) ListRoutingRules(ingressKeyID uint) []model.RoutingRule { c.mu.RLock() defer c.mu.RUnlock() out := c.routingRules[ingressKeyID] cp := make([]model.RoutingRule, len(out)) copy(cp, out) return cp } func (c *Store) SetRoutingRules(ingressKeyID uint, rules []model.RoutingRule) { c.mu.Lock() defer c.mu.Unlock() c.routingRules[ingressKeyID] = rules } func (c *Store) DeleteRoutingRulesForIngress(ingressKeyID uint) { c.mu.Lock() defer c.mu.Unlock() delete(c.routingRules, ingressKeyID) } type Counter struct { v atomic.Int64 } func (c *Counter) Add(n int64) { c.v.Add(n) } func (c *Counter) Load() int64 { return c.v.Load() }