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