busniess-user-center/pkg/middleware/jwt/jwt.go

package jwt

import (
	"fmt"
	"os"
	"time"

	contextUtil "busniess-user-center/pkg/utils/context"
	"busniess-user-center/pkg/utils/session"
	"busniess-user-center/pkg/utils/token"

	jwt "github.com/appleboy/gin-jwt/v2"
	"github.com/gin-gonic/gin"
	"github.com/spf13/cast"
	"go.uber.org/zap"
)

var (
	isLocalDevelopment = false
)

func init() {
	isLocalDevelopment = len(os.Getenv("LOCAL_DEVELOPMENT")) > 0
}

var (
	JwtIdentityKey = "id"
	Account        = "account"
	tokenName      = "token"
	AppCode        = "app_code"

	whiteUrlList = map[string]bool{
		"/api/user/login": true,
		"/sso":            true,
	}
)

type login struct {
	Account  string `form:"account" json:"account" binding:"required"`
	Password string `form:"password" json:"password" binding:"required"`
}

type JwtAuthMiddleware struct {
	Middleware  *jwt.GinJWTMiddleware
	authHandler gin.HandlerFunc
}

func NewJwtAuthMiddleware(logger *zap.SugaredLogger, refresher *token.TokenRefresher, secret []byte, timeout time.Duration) (*JwtAuthMiddleware, error) {
	authMiddleware, err := jwt.New(&jwt.GinJWTMiddleware{
		Realm:       "",
		Key:         secret,
		Timeout:     timeout,
		MaxRefresh:  time.Hour,
		IdentityKey: JwtIdentityKey,
		PayloadFunc: func(data interface{}) jwt.MapClaims {
			if v, ok := data.(*session.Session); ok {
				return jwt.MapClaims{
					JwtIdentityKey: v.ID,
					Account:        v.Account,
				}
			}
			return jwt.MapClaims{}
		},
		IdentityHandler: func(c *gin.Context) interface{} {
			var (
				id      uint = 0
				account      = ""
				appCode      = ""
			)

			claims := jwt.ExtractClaims(c)

			AccountInterface, ok := claims[Account]
			if ok {
				account = AccountInterface.(string)
			}

			idInterface, ok := claims[JwtIdentityKey]
			if ok {
				id = cast.ToUint(idInterface)
			}

			appCodeInterface, ok := claims[AppCode]
			if ok {
				appCode = appCodeInterface.(string)
			}

			// 从redis获取
			_, err := refresher.GetUserToken(id)
			if err != nil {
				return fmt.Errorf("没有权限")
			}

			return &session.Session{
				ID:      id,
				Account: account,
				AppCode: appCode,
			}
		},
		Authenticator: func(c *gin.Context) (interface{}, error) {
			var loginVals login
			if err := c.ShouldBind(&loginVals); err != nil {
				return "", jwt.ErrMissingLoginValues
			}
			userID := loginVals.Account
			password := loginVals.Password

			if userID == "admin" && password == "admin" {
				return &session.Session{
					ID:      1,
					Account: "admin",
				}, nil
			}

			return nil, jwt.ErrFailedAuthentication
		},
		Authorizator: func(data interface{}, c *gin.Context) bool {
			if v, ok := data.(*session.Session); ok {
				_ = contextUtil.PutSession(c, v)
				c.Request = c.Request.WithContext(contextUtil.PutSession(c.Request.Context(), v))
				return v.ID != 0
			}
			return false
		},
		Unauthorized: func(c *gin.Context, code int, message string) {
			if msg, ok := c.Get(JwtIdentityKey); ok {
				if err, ok := msg.(error); ok {
					message = err.Error()
				}
			}
			c.String(code, message)
		},

		TokenLookup:   "header:Authorization, cookie:" + tokenName,
		TokenHeadName: "Bearer",

		// TimeFunc provides the current time. You can override it to use another time value. This is useful for testing or if your server uses a different time zone than your tokens.
		TimeFunc: time.Now,
	})

	return &JwtAuthMiddleware{
		Middleware:  authMiddleware,
		authHandler: authMiddleware.MiddlewareFunc(),
	}, err
}

func (s *JwtAuthMiddleware) AuthHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		rPath := c.Request.URL.Path
		if _, ok := whiteUrlList[rPath]; !ok {
			s.authHandler(c)
		} else {
			c.Next()
		}
	}
}