llama-link/internal/api/middleware/auth.go

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)
}