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 == "" {
			allowedPaths := []string{"/api/onlyoffice/callback", "/api/files/download", "/api/files/thumbnail"}
			path := c.Path()
			for _, p := range allowedPaths {
				if strings.HasPrefix(path, p) {
					tokenString = c.Query("token")
					if tokenString != "" {
						isQueryToken = true
					}
					break
				}
			}
		}

		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 c.Path() == "/api/auth/refresh" {
				if tokenType != "refresh" {
					return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid token type for refresh endpoint"})
				}
			} 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, rememberMe bool) (string, error) {
	exp := time.Now().Add(7 * 24 * time.Hour).Unix() // 7 days
	if rememberMe {
		exp = time.Now().Add(21 * 24 * time.Hour).Unix() // 21 days (3 weeks)
	}

	claims := jwt.MapClaims{
		"sub":      userID,
		"username": username,
		"exp":      exp,
		"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, filePath string) (string, error) {
	claims := jwt.MapClaims{
		"sub":      userID,
		"username": username,
		"exp":      time.Now().Add(5 * time.Minute).Unix(),
		"iat":      time.Now().Unix(),
		"type":     "download",
	}

	if filePath != "" {
		claims["path"] = filePath
	}

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

// ValidateDownloadToken validates a download token
func ValidateDownloadToken(tokenString, secret, expectedUsername, expectedPath 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")
	}

	if expectedPath != "" {
		tokenPath, _ := claims["path"].(string)
		if tokenPath != "" && !strings.HasPrefix(expectedPath, tokenPath) {
			return fiber.NewError(fiber.StatusUnauthorized, "token not valid for this path")
		}
	}

	return nil
}