| package internal
|
|
|
| import (
|
| "crypto/rand"
|
| "crypto/subtle"
|
| "encoding/hex"
|
| "encoding/json"
|
| "fmt"
|
| "os"
|
| "path/filepath"
|
| "sort"
|
| "sync"
|
| "sync/atomic"
|
| "time"
|
| )
|
|
|
|
|
| type ApiKey struct {
|
| Key string `json:"key"`
|
| Name string `json:"name"`
|
| CreatedAt int64 `json:"created_at"`
|
| LastUsed int64 `json:"last_used"`
|
| UseCount int64 `json:"use_count"`
|
| Enabled bool `json:"enabled"`
|
| }
|
|
|
|
|
| type ApiKeyManager struct {
|
| mu sync.RWMutex
|
| keys map[string]*ApiKey
|
| dataDir string
|
| filename string
|
| }
|
|
|
| var (
|
| apiKeyManager *ApiKeyManager
|
| apiKeyOnce sync.Once
|
| )
|
|
|
|
|
| func GetApiKeyManager() *ApiKeyManager {
|
| apiKeyOnce.Do(func() {
|
| apiKeyManager = &ApiKeyManager{
|
| keys: make(map[string]*ApiKey),
|
| dataDir: "data",
|
| filename: "api_keys.json",
|
| }
|
| if err := apiKeyManager.load(); err != nil {
|
| LogWarn("ApiKeyManager 加载失败: %v", err)
|
| }
|
| })
|
| return apiKeyManager
|
| }
|
|
|
| func (m *ApiKeyManager) path() string {
|
| return filepath.Join(m.dataDir, m.filename)
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) load() error {
|
| if err := os.MkdirAll(m.dataDir, 0755); err != nil {
|
| return err
|
| }
|
| data, err := os.ReadFile(m.path())
|
| if err != nil {
|
| if os.IsNotExist(err) {
|
| return nil
|
| }
|
| return err
|
| }
|
| var arr []*ApiKey
|
| if err := json.Unmarshal(data, &arr); err != nil {
|
| return fmt.Errorf("parse %s: %w", m.path(), err)
|
| }
|
| m.mu.Lock()
|
| defer m.mu.Unlock()
|
| for _, k := range arr {
|
| if k.Key != "" {
|
| m.keys[k.Key] = k
|
| }
|
| }
|
| LogInfo("ApiKeyManager 已加载 %d 个 API Key", len(m.keys))
|
| return nil
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) save() error {
|
| m.mu.RLock()
|
| arr := make([]*ApiKey, 0, len(m.keys))
|
| for _, k := range m.keys {
|
| copy := *k
|
| arr = append(arr, ©)
|
| }
|
| m.mu.RUnlock()
|
| sort.Slice(arr, func(i, j int) bool { return arr[i].CreatedAt < arr[j].CreatedAt })
|
|
|
| data, err := json.MarshalIndent(arr, "", " ")
|
| if err != nil {
|
| return err
|
| }
|
| if err := os.MkdirAll(m.dataDir, 0755); err != nil {
|
| return err
|
| }
|
| tmpFile := m.path() + ".tmp"
|
| if err := os.WriteFile(tmpFile, data, 0644); err != nil {
|
| return err
|
| }
|
| return os.Rename(tmpFile, m.path())
|
| }
|
|
|
|
|
| func generateApiKey() string {
|
| b := make([]byte, 24)
|
| if _, err := rand.Read(b); err != nil {
|
|
|
| return fmt.Sprintf("sk-%d-%d", time.Now().UnixNano(), atomic.AddInt64(&fallbackKeyCounter, 1))
|
| }
|
| return "sk-" + hex.EncodeToString(b)
|
| }
|
|
|
| var fallbackKeyCounter int64
|
|
|
|
|
| func (m *ApiKeyManager) Create(name string) (*ApiKey, error) {
|
| if name == "" {
|
| name = "未命名"
|
| }
|
| if len(name) > 64 {
|
| name = name[:64]
|
| }
|
| k := &ApiKey{
|
| Key: generateApiKey(),
|
| Name: name,
|
| CreatedAt: time.Now().Unix(),
|
| Enabled: true,
|
| }
|
| m.mu.Lock()
|
| m.keys[k.Key] = k
|
| m.mu.Unlock()
|
| if err := m.save(); err != nil {
|
|
|
| m.mu.Lock()
|
| delete(m.keys, k.Key)
|
| m.mu.Unlock()
|
| return nil, fmt.Errorf("写文件失败: %w", err)
|
| }
|
| return k, nil
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) Delete(key string) error {
|
| m.mu.Lock()
|
| if _, ok := m.keys[key]; !ok {
|
| m.mu.Unlock()
|
| return fmt.Errorf("key 不存在")
|
| }
|
| delete(m.keys, key)
|
| m.mu.Unlock()
|
| return m.save()
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) SetEnabled(key string, enabled bool) error {
|
| m.mu.Lock()
|
| k, ok := m.keys[key]
|
| if !ok {
|
| m.mu.Unlock()
|
| return fmt.Errorf("key 不存在")
|
| }
|
| k.Enabled = enabled
|
| m.mu.Unlock()
|
| return m.save()
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) Validate(key string) bool {
|
| m.mu.RLock()
|
| k, ok := m.keys[key]
|
| m.mu.RUnlock()
|
| if !ok || !k.Enabled {
|
| return false
|
| }
|
| atomic.AddInt64(&k.UseCount, 1)
|
| atomic.StoreInt64(&k.LastUsed, time.Now().Unix())
|
|
|
| return true
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) List() []*ApiKey {
|
| m.mu.RLock()
|
| defer m.mu.RUnlock()
|
| out := make([]*ApiKey, 0, len(m.keys))
|
| for _, k := range m.keys {
|
| copy := *k
|
| out = append(out, ©)
|
| }
|
| sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
|
| return out
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) PersistCounters() {
|
| if err := m.save(); err != nil {
|
| LogWarn("ApiKeyManager 持久化计数失败: %v", err)
|
| }
|
| }
|
|
|
|
|
| func (m *ApiKeyManager) StartPeriodicSave() {
|
| go func() {
|
| ticker := time.NewTicker(60 * time.Second)
|
| defer ticker.Stop()
|
| for range ticker.C {
|
| m.PersistCounters()
|
| }
|
| }()
|
| }
|
|
|
|
|
|
|
|
|
| func ValidateAnyApiKey(key string) bool {
|
| if Cfg.SkipAuthToken {
|
| return true
|
| }
|
|
|
| for _, t := range Cfg.AuthTokens {
|
| if subtle.ConstantTimeCompare([]byte(t), []byte(key)) == 1 {
|
| return true
|
| }
|
| }
|
|
|
| return GetApiKeyManager().Validate(key)
|
| }
|
|
|
