llama-link/internal/auth/service.go

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
}