package auth import ( "crypto/rand" "encoding/hex" "errors" "time" "github.com/google/uuid" "golang.org/x/crypto/bcrypt" "gorm.io/gorm" "github.com/llamalink/llamalink/internal/db" ) const ( TokenPrefix = "llmk_" TokenLen = 32 // 64 hex chars ) var ( ErrInvalidKey = errors.New("invalid API key") ErrKeyNotFound = errors.New("API key not found") ErrKeyRevoked = errors.New("API key has been revoked") ErrKeyExpired = errors.New("API key has expired") ErrInsufficientScope = errors.New("insufficient scope") ) type CreateKeyRequest struct { Name string Scopes []string TokensLimit *int WebhookURL *string OwnerLabel *string } type Service struct { db *gorm.DB } func NewService(db *gorm.DB) *Service { return &Service{db: db} } func (s *Service) GenerateToken() (string, error) { bytes := make([]byte, TokenLen) if _, err := rand.Read(bytes); err != nil { return "", err } return TokenPrefix + hex.EncodeToString(bytes), nil } func (s *Service) HashToken(token string) string { hash, _ := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost) return string(hash) } func (s *Service) Create(req CreateKeyRequest) (*db.ApiKey, string, error) { token, err := s.GenerateToken() if err != nil { return nil, "", err } isAdmin := s.isFirstKey() apiKey := &db.ApiKey{ ID: uuid.New(), Name: req.Name, KeyHash: s.HashToken(token), KeyPrefix: token[:len(TokenPrefix)+8], Scopes: req.Scopes, IsActive: true, IsAdmin: isAdmin, OwnerLabel: req.OwnerLabel, } if err := s.db.Create(apiKey).Error; err != nil { return nil, "", err } // Create quota if specified if req.TokensLimit != nil && *req.TokensLimit > 0 { now := time.Now().UTC() periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) periodEnd := periodStart.AddDate(0, 1, 0) quota := &db.Quota{ ID: uuid.New(), ApiKeyID: apiKey.ID, PeriodStart: periodStart, PeriodEnd: periodEnd, TokensLimit: *req.TokensLimit, TokensUsed: 0, } s.db.Create(quota) } return apiKey, token, nil } func (s *Service) Revoke(id uuid.UUID) error { result := s.db.Model(&db.ApiKey{}).Where("id = ?", id).Update("is_active", false) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { return ErrKeyNotFound } return nil } func (s *Service) Validate(token string) (*db.ApiKey, error) { if len(token) < len(TokenPrefix)+8 { return nil, ErrInvalidKey } prefix := token[:len(TokenPrefix)+8] var apiKey db.ApiKey if err := s.db.Where("key_prefix = ? AND is_active = ?", prefix, true).First(&apiKey).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrInvalidKey } return nil, err } if err := bcrypt.CompareHashAndPassword([]byte(apiKey.KeyHash), []byte(token)); err != nil { return nil, ErrInvalidKey } if apiKey.ExpiresAt != nil && time.Now().After(*apiKey.ExpiresAt) { return nil, ErrKeyExpired } // Update last used now := time.Now() s.db.Model(&apiKey).Update("last_used_at", now) return &apiKey, nil } func (s *Service) List(includeInactive bool) ([]db.ApiKey, error) { var keys []db.ApiKey query := s.db.Order("created_at DESC") if !includeInactive { query = query.Where("is_active = ?", true) } if err := query.Find(&keys).Error; err != nil { return nil, err } return keys, nil } func (s *Service) GetByID(id uuid.UUID) (*db.ApiKey, error) { var key db.ApiKey if err := s.db.Where("id = ?", id).First(&key).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrKeyNotFound } return nil, err } return &key, nil } func (s *Service) HasScope(key *db.ApiKey, scope string) bool { for _, kScope := range key.Scopes { if kScope == scope { return true } } return false } func (s *Service) isFirstKey() bool { var count int64 s.db.Model(&db.ApiKey{}).Count(&count) return count == 0 }