Godemo/web-training/internal/middleware/auth.go

package middleware

import (
	"net/http"
	"strings"
	"web-training/pkg/auth"
	"web-training/internal/utils"

	"github.com/gin-gonic/gin"
)

// AuthMiddleware JWT认证中间件
func AuthMiddleware(jwtService *auth.JWTService) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 从请求头获取Authorization
		authHeader := c.GetHeader("Authorization")
		if authHeader == "" {
			utils.ErrorResponse(c, http.StatusUnauthorized, "缺少认证令牌", "")
			c.Abort()
			return
		}

		// 检查Bearer前缀
		parts := strings.Split(authHeader, " ")
		if len(parts) != 2 || parts[0] != "Bearer" {
			utils.ErrorResponse(c, http.StatusUnauthorized, "令牌格式错误", "")
			c.Abort()
			return
		}

		// 验证令牌
		token := parts[1]
		claims, err := jwtService.ValidateToken(token)
		if err != nil {
			utils.ErrorResponse(c, http.StatusUnauthorized, "无效令牌", err.Error())
			c.Abort()
			return
		}

		// 将用户信息存储到上下文
		c.Set("userID", claims.UserID)
		c.Set("userRoles", claims.Roles)
		c.Set("claims", claims)

		c.Next()
	}
}

// RoleBasedAuthMiddleware 基于角色的权限控制中间件
func RoleBasedAuthMiddleware(requiredRoles ...string) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 获取用户角色
		userRoles, exists := c.Get("userRoles")
		if !exists {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "无法获取用户角色")
			c.Abort()
			return
		}

		// 检查用户角色是否包含所需角色
		userRolesSlice, ok := userRoles.([]string)
		if !ok {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "角色格式错误")
			c.Abort()
			return
		}

		hasPermission := false
		for _, requiredRole := range requiredRoles {
			for _, userRole := range userRolesSlice {
				if userRole == requiredRole {
					hasPermission = true
					break
				}
			}
			if hasPermission {
				break
			}
		}

		if !hasPermission {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "角色权限不足")
			c.Abort()
			return
		}

		c.Next()
	}
}

// AdminOnlyMiddleware 仅管理员访问中间件
func AdminOnlyMiddleware() gin.HandlerFunc {
	return RoleBasedAuthMiddleware("admin")
}

// SelfOrAdminMiddleware 用户本人或管理员访问中间件
func SelfOrAdminMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		userID, exists := c.Get("userID")
		if !exists {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "无法获取用户ID")
			c.Abort()
			return
		}

		userRoles, exists := c.Get("userRoles")
		if !exists {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "无法获取用户角色")
			c.Abort()
			return
		}

		// 获取请求中的用户ID
		paramID := c.Param("id")
		if paramID == "" {
			utils.ErrorResponse(c, http.StatusBadRequest, "请求参数错误", "缺少用户ID参数")
			c.Abort()
			return
		}

		// 检查是否是管理员
		userRolesSlice, ok := userRoles.([]string)
		if !ok {
			utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "角色格式错误")
			c.Abort()
			return
		}

		isAdmin := false
		for _, role := range userRolesSlice {
			if role == "admin" {
				isAdmin = true
				break
			}
		}

		// 如果不是管理员，检查是否是用户本人
		if !isAdmin {
			userIDUint, ok := userID.(uint)
			if !ok {
				utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "用户ID格式错误")
				c.Abort()
				return
			}

			if paramID != string(rune(userIDUint)) {
				utils.ErrorResponse(c, http.StatusForbidden, "权限不足", "只能访问自己的资源")
				c.Abort()
				return
			}
		}

		c.Next()
	}
}