tm_back/internal/user/middleware.go

package user

import (
	"net/http"
	"strings"
	"tm/pkg/audit"
	"tm/pkg/response"

	"github.com/labstack/echo/v4"
	"go.mongodb.org/mongo-driver/v2/bson"
)

// AuthMiddleware validates JWT access tokens and extracts user information
func (h *Handler) AuthMiddleware() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Extract token from Authorization header
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader == "" {
				return response.Unauthorized(c, "Authorization header required")
			}

			// Check if it's a Bearer token
			if !strings.HasPrefix(authHeader, "Bearer ") {
				return response.Unauthorized(c, "Invalid authorization format. Use 'Bearer <token>'")
			}

			// Extract the token
			tokenString := strings.TrimPrefix(authHeader, "Bearer ")

			// Validate the access token
			validationResult, err := h.authService.ValidateAccessToken(tokenString)
			if err != nil {
				h.logger.Error("Token validation error", map[string]interface{}{
					"error": err.Error(),
				})
				return response.Unauthorized(c, "Invalid token")
			}

			if !validationResult.Valid {
				if validationResult.Expired {
					return response.Unauthorized(c, "Token expired")
				}
				return response.Unauthorized(c, validationResult.Error)
			}

			// Extract user information from token
			userID, err := bson.ObjectIDFromHex(validationResult.Claims.UserID)
			if err != nil {
				h.logger.Error("Failed to parse user ID from token", map[string]interface{}{
					"error": err.Error(),
				})
				return response.Unauthorized(c, "Invalid token format")
			}

			// Store user information in context for handlers to use
			c.Set("user_id", userID.Hex())
			c.Set("user_email", validationResult.Claims.Email)
			c.Set("user_role", validationResult.Claims.Role)
			c.Set("company_id", validationResult.Claims.CompanyID)
			c.Set("access_token", tokenString)

			audit.EnrichEchoContext(c, audit.ActorTypeAdmin)

			// Continue to next handler
			return next(c)
		}
	}
}

// AdminMiddleware checks if the authenticated user has admin role
func (h *Handler) AdminMiddleware() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Get user role from context (set by AuthMiddleware)
			userRole, ok := c.Get("user_role").(string)
			if !ok {
				return response.Unauthorized(c, "User role not found in context")
			}

			// Check if user has admin role
			if userRole != string(UserRoleAdmin) {
				return response.Forbidden(c, "Admin access required")
			}

			// Continue to next handler
			return next(c)
		}
	}
}

// RoleMiddleware checks if the authenticated user has the required role
func (h *Handler) RoleMiddleware(requiredRole UserRole) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Get user role from context (set by AuthMiddleware)
			userRole, ok := c.Get("user_role").(string)
			if !ok {
				return response.Unauthorized(c, "User role not found in context")
			}

			// Check if user has the required role
			if userRole != string(requiredRole) {
				return response.Forbidden(c, "Insufficient permissions")
			}

			// Continue to next handler
			return next(c)
		}
	}
}

// CompanyAccessMiddleware checks if the authenticated user has access to the specified company
func (h *Handler) CompanyAccessMiddleware() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Get user information from context
			userRole, ok := c.Get("user_role").(string)
			if !ok {
				return response.Unauthorized(c, "User role not found in context")
			}

			// Admin users have access to all companies
			if userRole == string(UserRoleAdmin) {
				return next(c)
			}

			// Get company ID from context
			userCompanyID, ok := c.Get("company_id").(string)
			if !ok || userCompanyID == "" {
				return response.Forbidden(c, "Company access not available")
			}

			// Get company ID from URL parameter
			requestedCompanyID := c.Param("company_id")
			if requestedCompanyID == "" {
				// If no company ID in URL, allow access to user's own company
				return next(c)
			}

			// Check if user has access to the requested company
			if userCompanyID != requestedCompanyID {
				return response.Forbidden(c, "Access to this company not allowed")
			}

			// Continue to next handler
			return next(c)
		}
	}
}

// GetUserIDFromContext extracts user ID from Echo context
func GetUserIDFromContext(c echo.Context) (string, error) {
	userID, ok := c.Get("user_id").(string)
	if !ok {
		return "", echo.NewHTTPError(http.StatusUnauthorized, "User ID not found in context")
	}
	return userID, nil
}

// GetUserEmailFromContext extracts user email from Echo context
func GetUserEmailFromContext(c echo.Context) (string, error) {
	userEmail, ok := c.Get("user_email").(string)
	if !ok {
		return "", echo.NewHTTPError(http.StatusUnauthorized, "User email not found in context")
	}
	return userEmail, nil
}

// GetUserRoleFromContext extracts user role from Echo context
func GetUserRoleFromContext(c echo.Context) (string, error) {
	userRole, ok := c.Get("user_role").(string)
	if !ok {
		return "", echo.NewHTTPError(http.StatusUnauthorized, "User role not found in context")
	}
	return userRole, nil
}

// GetCompanyIDFromContext extracts company ID from Echo context
func GetCompanyIDFromContext(c echo.Context) (string, error) {
	companyID, ok := c.Get("company_id").(string)
	if !ok {
		return "", echo.NewHTTPError(http.StatusUnauthorized, "Company ID not found in context")
	}
	return companyID, nil
}

// GetAccessTokenFromContext extracts access token from Echo context
func GetAccessTokenFromContext(c echo.Context) (string, error) {
	accessToken, ok := c.Get("access_token").(string)
	if !ok {
		return "", echo.NewHTTPError(http.StatusUnauthorized, "Access token not found in context")
	}
	return accessToken, nil
}