package middleware import ( "net/http" "strings" "time" "stock-user-system/internal/config" "stock-user-system/internal/models" "github.com/gin-gonic/gin" "github.com/golang-jwt/jwt/v5" "gorm.io/gorm" ) type CurrentUser struct { ID string Username string Role models.RoleName } var publicPaths = map[string]struct{}{ "/api/auth/register": {}, "/api/auth/login": {}, } func AuthMiddleware(cfg *config.Config, db *gorm.DB) gin.HandlerFunc { return func(c *gin.Context) { authHeader := c.GetHeader("Authorization") token := "" if strings.HasPrefix(authHeader, "Bearer ") { token = strings.TrimPrefix(authHeader, "Bearer ") } if token != "" { claims, err := parseToken(token, cfg.JWTSecret) if err == nil { var user models.User if err := db.Preload("Role").Where("id = ? AND status = ?", claims.Subject, "active").First(&user).Error; err == nil { c.Set("currentUser", CurrentUser{ ID: user.ID, Username: user.Username, Role: user.Role.NameEnum(), }) } } } if _, exists := c.Get("currentUser"); !exists { if c.Request.Method == "OPTIONS" { c.Next() return } if _, ok := publicPaths[c.Request.URL.Path]; !ok { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "error": "访问令牌无效或已过期"}) return } } c.Next() } } type Claims struct { Subject string `json:"sub"` Role string `json:"role"` jwt.RegisteredClaims } func GenerateToken(userID string, role models.RoleName, cfg *config.Config) (string, error) { claims := Claims{ Subject: userID, Role: role.String(), RegisteredClaims: jwt.RegisteredClaims{ ExpiresAt: jwt.NewNumericDate(time.Now().Add(cfg.JWTExpiration())), }, } token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims) return token.SignedString([]byte(cfg.JWTSecret)) } func parseToken(tokenString string, secret string) (*Claims, error) { token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) { return []byte(secret), nil }) if err != nil { return nil, err } if claims, ok := token.Claims.(*Claims); ok && token.Valid { return claims, nil } return nil, jwt.ErrSignatureInvalid } func RequireAuth() gin.HandlerFunc { return func(c *gin.Context) { if _, exists := c.Get("currentUser"); !exists { c.JSON(http.StatusUnauthorized, gin.H{"success": false, "error": "访问令牌无效或已过期"}) c.Abort() return } c.Next() } } func RequireRoles(allowed ...models.RoleName) gin.HandlerFunc { return func(c *gin.Context) { val, exists := c.Get("currentUser") if !exists { c.JSON(http.StatusUnauthorized, gin.H{"success": false, "error": "访问令牌无效或已过期"}) c.Abort() return } current := val.(CurrentUser) for _, role := range allowed { if current.Role == role { c.Next() return } } c.JSON(http.StatusForbidden, gin.H{"success": false, "error": "权限不足"}) c.Abort() } } func GetCurrentUser(c *gin.Context) (CurrentUser, bool) { val, exists := c.Get("currentUser") if !exists { return CurrentUser{}, false } return val.(CurrentUser), true }