File size: 3,290 Bytes
f1dd159 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | package clientkey
import (
"sync"
"time"
clientkeydomain "github.com/chenyme/grok2api/backend/internal/domain/clientkey"
)
const (
keyTouchInterval = time.Minute
touchTrackerMaxEntries = 10000
keyAuthCacheTTL = time.Second
keyAuthCacheMaxEntries = 10000
)
type cachedAuthKey struct {
value clientkeydomain.Key
expiresAt time.Time
}
type authKeyCache struct {
mu sync.RWMutex
byPrefix map[string]cachedAuthKey
}
func newAuthKeyCache() *authKeyCache {
return &authKeyCache{byPrefix: make(map[string]cachedAuthKey)}
}
func (c *authKeyCache) get(prefix string, now time.Time) (clientkeydomain.Key, bool) {
c.mu.RLock()
entry, ok := c.byPrefix[prefix]
c.mu.RUnlock()
if !ok || !now.Before(entry.expiresAt) {
if ok {
c.mu.Lock()
delete(c.byPrefix, prefix)
c.mu.Unlock()
}
return clientkeydomain.Key{}, false
}
value := entry.value
value.AllowedModels = append([]uint64(nil), entry.value.AllowedModels...)
return value, true
}
func (c *authKeyCache) put(prefix string, value clientkeydomain.Key, now time.Time) {
if prefix == "" || value.BillingLimitUSDTicks > 0 {
return
}
value.EncryptedSecret = ""
value.AllowedModels = append([]uint64(nil), value.AllowedModels...)
c.mu.Lock()
defer c.mu.Unlock()
c.byPrefix[prefix] = cachedAuthKey{value: value, expiresAt: now.Add(keyAuthCacheTTL)}
if len(c.byPrefix) <= keyAuthCacheMaxEntries {
return
}
for candidate, entry := range c.byPrefix {
if !now.Before(entry.expiresAt) {
delete(c.byPrefix, candidate)
}
}
for len(c.byPrefix) > keyAuthCacheMaxEntries {
for candidate := range c.byPrefix {
delete(c.byPrefix, candidate)
break
}
}
}
func (c *authKeyCache) deleteID(id uint64) {
c.mu.Lock()
defer c.mu.Unlock()
for prefix, entry := range c.byPrefix {
if entry.value.ID == id {
delete(c.byPrefix, prefix)
}
}
}
func (c *authKeyCache) deleteIDs(ids []uint64) {
set := make(map[uint64]struct{}, len(ids))
for _, id := range ids {
set[id] = struct{}{}
}
c.mu.Lock()
defer c.mu.Unlock()
for prefix, entry := range c.byPrefix {
if _, ok := set[entry.value.ID]; ok {
delete(c.byPrefix, prefix)
}
}
}
func (c *authKeyCache) clear() {
c.mu.Lock()
clear(c.byPrefix)
c.mu.Unlock()
}
// touchTracker 合并非关键的最近使用时间写入。
type touchTracker struct {
mu sync.Mutex
lastTouched map[uint64]time.Time
}
func newTouchTracker() *touchTracker {
return &touchTracker{lastTouched: make(map[uint64]time.Time)}
}
func (c *touchTracker) deleteID(id uint64) {
c.mu.Lock()
defer c.mu.Unlock()
delete(c.lastTouched, id)
}
func (c *touchTracker) deleteIDs(ids []uint64) {
c.mu.Lock()
defer c.mu.Unlock()
for _, id := range ids {
delete(c.lastTouched, id)
}
}
func (c *touchTracker) shouldTouch(id uint64, now time.Time) bool {
c.mu.Lock()
defer c.mu.Unlock()
if last := c.lastTouched[id]; !last.IsZero() && now.Sub(last) < keyTouchInterval {
return false
}
c.lastTouched[id] = now
if len(c.lastTouched) > touchTrackerMaxEntries {
var oldestID uint64
var oldest time.Time
for candidateID, touchedAt := range c.lastTouched {
if oldestID == 0 || touchedAt.Before(oldest) {
oldestID = candidateID
oldest = touchedAt
}
}
delete(c.lastTouched, oldestID)
}
return true
}
|