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" AdminCtx = "admin_user" ) 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 AdminSessionAuth(jwtService *auth.AdminJWTService) gin.HandlerFunc { return func(c *gin.Context) { authHeader := c.GetHeader("Authorization") if authHeader == "" { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Authorization header required", }, }) return } token := strings.TrimPrefix(authHeader, "Bearer ") if token == authHeader || strings.HasPrefix(token, auth.TokenPrefix) { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Admin session required", }, }) return } claims, err := jwtService.VerifyToken(token) if err != nil { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "error": gin.H{ "code": "unauthorized", "message": "Invalid or expired session", }, }) return } c.Set(AdminCtx, claims) 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) } func GetAdminClaims(c *gin.Context) *auth.AdminClaims { claims, _ := c.Get(AdminCtx) if claims == nil { return nil } return claims.(*auth.AdminClaims) }