Yimaru-BackEnd/internal/web_server/middleware.go

package httpserver

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/SamuelTariku/FortuneBet-Backend/internal/domain"
	jwtutil "github.com/SamuelTariku/FortuneBet-Backend/internal/web_server/jwt"
	"github.com/gofiber/fiber/v2"
	"go.uber.org/zap"
)

func (a *App) authMiddleware(c *fiber.Ctx) error {
	ip := c.IP()
	userAgent := c.Get("User-Agent")
	c.Locals("ip_address", ip)
	c.Locals("user_agent", userAgent)
	authHeader := c.Get("Authorization")
	if authHeader == "" {
		a.mongoLoggerSvc.Info("Authorization header missing",
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Authorization header missing")
	}

	if !strings.HasPrefix(authHeader, "Bearer ") {
		a.mongoLoggerSvc.Info("Invalid authorization header format",
			zap.String("authHeader", authHeader),
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Invalid authorization header format")
	}

	accessToken := strings.TrimPrefix(authHeader, "Bearer ")
	c.Locals("access_token", accessToken)
	claim, err := jwtutil.ParseJwt(accessToken, a.JwtConfig.JwtAccessKey)
	if err != nil {
		if errors.Is(err, jwtutil.ErrExpiredToken) {
			a.mongoLoggerSvc.Info("Access Token Expired",
				zap.Int("status_code", fiber.StatusUnauthorized),
				zap.String("ip_address", ip),
				zap.String("user_agent", userAgent),
				zap.Time("timestamp", time.Now()),
			)
			return fiber.NewError(fiber.StatusUnauthorized, "Access token expired")
		}
		a.mongoLoggerSvc.Info("Invalid Access Token",
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Invalid access token")
	}

	refreshToken := c.Get("Refresh-Token")
	if refreshToken == "" {

		// refreshToken = c.Cookies("refresh_token", "")
		// return fiber.NewError(fiber.StatusUnauthorized, "Refresh token missing")

	}
	// Asserting to make sure that there is no company role without a valid company id
	if claim.Role != domain.RoleSuperAdmin && claim.Role != domain.RoleCustomer && !claim.CompanyID.Valid {
		a.mongoLoggerSvc.Error("Company Role without Company ID",
			zap.Int64("userID", claim.UserId),
			zap.Int("status_code", fiber.StatusInternalServerError),
			zap.Error(err),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusInternalServerError, "Company Role without Company ID")
	}
	c.Locals("user_id", claim.UserId)
	c.Locals("role", claim.Role)
	c.Locals("company_id", claim.CompanyID)
	c.Locals("refresh_token", refreshToken)

	var branchID domain.ValidInt64

	if claim.Role == domain.RoleCashier {
		branch, err := a.branchSvc.GetBranchByCashier(c.Context(), claim.UserId)
		if err != nil {
			a.mongoLoggerSvc.Error("Failed to get branch id for cashier",
				zap.Int64("userID", claim.UserId),
				zap.Int("status_code", fiber.StatusInternalServerError),
				zap.Error(err),
				zap.Time("timestamp", time.Now()),
			)
			return fiber.NewError(fiber.StatusInternalServerError, "Failed to branch id for cashier")
		}
		branchID = domain.ValidInt64{
			Value: branch.ID,
			Valid: true,
		}

	}

	c.Locals("branch_id", branchID)
	return c.Next()
}

func (a *App) SuperAdminOnly(c *fiber.Ctx) error {
	userID := c.Locals("user_id").(int64)
	userRole := c.Locals("role").(domain.Role)
	if userRole != domain.RoleSuperAdmin {
		a.mongoLoggerSvc.Warn("Attempt to access restricted SuperAdminOnly route",
			zap.Int64("userID", userID),
			zap.String("role", string(userRole)),
			zap.Int("status_code", fiber.StatusForbidden),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusForbidden, "This route is restricted")
	}
	return c.Next()
}

func (a *App) CompanyOnly(c *fiber.Ctx) error {
	userID := c.Locals("user_id").(int64)
	userRole := c.Locals("role").(domain.Role)
	if userRole == domain.RoleCustomer {
		a.mongoLoggerSvc.Warn("Attempt to access restricted CompanyOnly route",
			zap.Int64("userID", userID),
			zap.String("role", string(userRole)),
			zap.Int("status_code", fiber.StatusForbidden),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusForbidden, "This route is restricted")
	}
	return c.Next()
}

func (a *App) OnlyAdminAndAbove(c *fiber.Ctx) error {
	userID := c.Locals("user_id").(int64)
	userRole := c.Locals("role").(domain.Role)
	if userRole != domain.RoleSuperAdmin && userRole != domain.RoleAdmin {
		a.mongoLoggerSvc.Warn("Attempt to access restricted OnlyAdminAndAbove route",
			zap.Int64("userID", userID),
			zap.String("role", string(userRole)),
			zap.Int("status_code", fiber.StatusForbidden),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusForbidden, "This route is restricted")
	}
	return c.Next()
}

func (a *App) OnlyBranchManagerAndAbove(c *fiber.Ctx) error {
	userID := c.Locals("user_id").(int64)
	userRole := c.Locals("role").(domain.Role)
	if userRole != domain.RoleSuperAdmin && userRole != domain.RoleAdmin && userRole != domain.RoleBranchManager {
		a.mongoLoggerSvc.Warn("Attempt to access restricted OnlyBranchMangerAndAbove route",
			zap.Int64("userID", userID),
			zap.String("role", string(userRole)),
			zap.Int("status_code", fiber.StatusForbidden),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusForbidden, "This route is restricted")
	}
	return c.Next()
}

func (a *App) WebsocketAuthMiddleware(c *fiber.Ctx) error {
	tokenStr := c.Query("token")
	ip := c.IP()
	userAgent := c.Get("User-Agent")

	if tokenStr == "" {
		a.mongoLoggerSvc.Info("Missing token in query parameter",
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Missing token")
	}

	claim, err := jwtutil.ParseJwt(tokenStr, a.JwtConfig.JwtAccessKey)
	if err != nil {
		if errors.Is(err, jwtutil.ErrExpiredToken) {
			a.mongoLoggerSvc.Info("Token expired",
				zap.Int("status_code", fiber.StatusUnauthorized),
				zap.String("ip_address", ip),
				zap.String("user_agent", userAgent),
				zap.Time("timestamp", time.Now()),
			)
			return fiber.NewError(fiber.StatusUnauthorized, "Token expired")
		}
		a.logger.Error("Invalid token", "error", err)
		a.mongoLoggerSvc.Info("Invalid token",
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Invalid token")
	}

	userID := claim.UserId
	if userID == 0 {
		a.mongoLoggerSvc.Info("Invalid user ID in token claims",
			zap.Int("status_code", fiber.StatusUnauthorized),
			zap.String("ip_address", ip),
			zap.String("user_agent", userAgent),
			zap.Time("timestamp", time.Now()),
		)
		return fiber.NewError(fiber.StatusUnauthorized, "Invalid user ID")
	}

	c.Locals("userID", userID)
	a.mongoLoggerSvc.Info("Authenticated WebSocket connection",
		zap.Int64("userID", userID),
		zap.Time("timestamp", time.Now()),
	)
	return c.Next()
}