drive/backend/middleware/auth.go

package middleware

import (
	"strings"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/golang-jwt/jwt/v5"
)

// AuthMiddleware validates JWT tokens from the Authorization header.
func AuthMiddleware(jwtSecret string) fiber.Handler {
	return func(c *fiber.Ctx) error {
		authHeader := c.Get("Authorization")
		tokenString := ""

		if authHeader != "" {
			parts := strings.SplitN(authHeader, " ", 2)
			if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {
				tokenString = parts[1]
			}
		}

		isQueryToken := false
		if tokenString == "" {
			tokenString = c.Query("token")
			isQueryToken = true
		}

		if tokenString == "" {
			return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{
				"error": "missing or invalid authorization",
			})
		}

		token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fiber.NewError(fiber.StatusUnauthorized, "unexpected signing method")
			}
			return []byte(jwtSecret), nil
		})

		if err != nil || !token.Valid {
			return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{
				"error": "invalid or expired token",
			})
		}

		claims, ok := token.Claims.(jwt.MapClaims)
		if !ok {
			return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{
				"error": "invalid token claims",
			})
		}

		tokenType, _ := claims["type"].(string)
		if isQueryToken {
			if tokenType != "download" {
				return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid token type for query parameter"})
			}
		} else {
			if tokenType != "access" {
				return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid token type for authorization header"})
			}
		}

		c.Locals("userID", claims["sub"])
		c.Locals("username", claims["username"])

		return c.Next()
	}
}

// GenerateAccessToken creates a short-lived JWT access token.
func GenerateAccessToken(userID int, username, secret string, expirySecs int) (string, error) {
	claims := jwt.MapClaims{
		"sub":      userID,
		"username": username,
		"exp":      time.Now().Add(time.Duration(expirySecs) * time.Second).Unix(),
		"iat":      time.Now().Unix(),
		"type":     "access",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(secret))
}

// GenerateRefreshToken creates a longer-lived refresh token.
func GenerateRefreshToken(userID int, username, secret string) (string, error) {
	claims := jwt.MapClaims{
		"sub":      userID,
		"username": username,
		"exp":      time.Now().Add(7 * 24 * time.Hour).Unix(), // 7 days
		"iat":      time.Now().Unix(),
		"type":     "refresh",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(secret))
}

// GenerateDownloadToken creates a short-lived download token.
func GenerateDownloadToken(userID int, username, secret string) (string, error) {
	claims := jwt.MapClaims{
		"sub":      userID,
		"username": username,
		"exp":      time.Now().Add(5 * time.Minute).Unix(),
		"iat":      time.Now().Unix(),
		"type":     "download",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(secret))
}

// ValidateDownloadToken validates a download token
func ValidateDownloadToken(tokenString, secret, expectedUsername string) error {
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		return []byte(secret), nil
	})

	if err != nil || !token.Valid {
		return fiber.NewError(fiber.StatusUnauthorized, "invalid token")
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok || claims["type"] != "download" {
		return fiber.NewError(fiber.StatusUnauthorized, "invalid token claims")
	}

	if expectedUsername != "" && claims["username"] != expectedUsername {
		return fiber.NewError(fiber.StatusUnauthorized, "token belongs to different context")
	}

	return nil
}