package middleware import ( "net/http" "strings" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/llamalink/llamalink/internal/auth" "github.com/llamalink/llamalink/internal/db" ) const ( ApiKeyCtx = "api_key" ApiKeyIDCtx = "api_key_id" ) func APIKeyAuth(authService *auth.Service) gin.HandlerFunc { return func(c *gin.Context) { authHeader := c.GetHeader("Authorization") if authHeader == "" { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "invalid_api_key", "message": "Authorization header required", }, }) return } token := strings.TrimPrefix(authHeader, "Bearer ") if token == authHeader { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "invalid_api_key", "message": "Bearer token required", }, }) return } apiKey, err := authService.Validate(token) if err != nil { code := "invalid_api_key" status := http.StatusUnauthorized if err == auth.ErrKeyRevoked || err == auth.ErrKeyExpired { code = "api_key_revoked" status = http.StatusUnauthorized } c.AbortWithStatusJSON(status, gin.H{ "error": gin.H{ "code": code, "message": err.Error(), }, }) return } c.Set(ApiKeyCtx, apiKey) c.Set(ApiKeyIDCtx, apiKey.ID) c.Next() } } func RequireScope(authService *auth.Service, scope string) gin.HandlerFunc { return func(c *gin.Context) { apiKey, exists := c.Get(ApiKeyCtx) if !exists { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "invalid_api_key", "message": "Authentication required", }, }) return } key := apiKey.(*db.ApiKey) if !authService.HasScope(key, scope) { c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ "error": gin.H{ "code": "insufficient_scope", "message": "API key lacks required scope: " + scope, }, }) return } c.Next() } } func AdminOnly() gin.HandlerFunc { return func(c *gin.Context) { apiKey, exists := c.Get(ApiKeyCtx) if !exists { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Admin access required", }, }) return } key := apiKey.(*db.ApiKey) if !key.IsAdmin { c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Admin access required", }, }) return } c.Next() } } func AdminTokenAuth(token string) gin.HandlerFunc { return func(c *gin.Context) { adminToken := c.GetHeader("X-Admin-Token") if adminToken == "" { adminToken = c.GetHeader("Authorization") adminToken = strings.TrimPrefix(adminToken, "Bearer ") } if adminToken != token { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Invalid admin token", }, }) return } c.Next() } } func GetAPIKeyID(c *gin.Context) uuid.UUID { id, _ := c.Get(ApiKeyIDCtx) return id.(uuid.UUID) } func GetAPIKey(c *gin.Context) *db.ApiKey { key, _ := c.Get(ApiKeyCtx) if key == nil { return nil } return key.(*db.ApiKey) }