moutai/src/api/jwt.go

package api

import (
	"fmt"
	"log"
	"moutai/conf"
	"time"

	jwt "github.com/appleboy/gin-jwt/v2"
	"github.com/gin-gonic/gin"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

var identityKey = "id"

type UtilsJwt struct {
	jwt *jwt.GinJWTMiddleware
}

func (j *UtilsJwt) MiddleFunc() gin.HandlerFunc {
	return j.jwt.MiddlewareFunc()
}

func (j *UtilsJwt) JwtCreateToken(data interface{}) (string, time.Time, error) {
	return j.jwt.TokenGenerator(data)
}

func NewUitlsJwt(app *conf.AppConf) *UtilsJwt {

	var utils = &UtilsJwt{
		jwt: nil,
	}

	var jwtImpl *jwt.GinJWTMiddleware

	jwtImpl, err := jwt.New(&jwt.GinJWTMiddleware{
		Realm:       app.Jwt.Realm,
		Key:         []byte(app.Jwt.Key),
		Timeout:     time.Hour * time.Duration(app.Jwt.TimeoutHour),
		MaxRefresh:  time.Hour * time.Duration(app.Jwt.TimeoutHour),
		IdentityKey: identityKey,
		PayloadFunc: func(data interface{}) jwt.MapClaims {
			if v, ok := data.(*JWTUser); ok {
				return jwt.MapClaims{
					"id": v.ID,
				}
			}
			return jwt.MapClaims{}
		},
		IdentityHandler: func(c *gin.Context) interface{} {
			claims := jwt.ExtractClaims(c)
			// lg.Debug().Msgf("token: %v\n", claims)

			u := &JWTUser{
				ID: claims["id"].(string), //uint32(claims["id"].(float64)),
			}

			return u
		},

		Authenticator: func(c *gin.Context) (interface{}, error) {
			return &JWTUser{ID: ""}, nil
		},

		Authorizator: func(data interface{}, _ *gin.Context) bool {
			u := data.(*JWTUser)
			uid, _ := primitive.ObjectIDFromHex(u.ID)

			return !uid.IsZero()
		},

		LoginResponse: func(c *gin.Context, status int, token string, expire time.Time) {

			t, _ := jwtImpl.ParseTokenString(token)
			fmt.Println("LoginResponse==>", status, token)

			ResultSuc(c, map[string]interface{}{
				"user":   t.Claims,
				"token":  token,
				"expire": expire.Format(time.RFC3339),
			})
		},

		LogoutResponse: func(c *gin.Context, _ int) {
			ResultSuc(c, true)
		},

		RefreshResponse: func(c *gin.Context, status int, token string, expire time.Time) {

			ResultSuc(c, map[string]interface{}{
				"status": status,
				"token":  token,
				"expire": expire.Format(time.RFC3339),
			})
		},

		Unauthorized: func(c *gin.Context, status int, token string) {
			if token == "用户名/密码 不正确" {
				ResultFail401(c, "账号密码不对", map[string]interface{}{
					"status": -1,
					"token":  token,
				})
				return
			}

			fmt.Println("xxxxx")
			fmt.Println(token, status)

			ResultFail401(c, token, map[string]interface{}{
				"status": status,
				"token":  token,
			})
		},

		HTTPStatusMessageFunc: func(e error, _ *gin.Context) string {

			// if e == jwt.ErrFailedAuthentication {
			// 	return "用户名/密码 不正确"
			// }
			fmt.Println("HTTPStatusMessageFunc", e)
			return e.Error()
		},

		// TokenLookup is a string in the form of "<source>:<name>" that is used
		// to extract token from the request.
		// Optional. Default value "header:Authorization".
		// Possible values:
		// - "header:<name>"
		// - "query:<name>"
		// - "cookie:<name>"
		// - "param:<name>"
		TokenLookup: "header: Authorization, query: token, cookie: jwt",
		// TokenLookup: "query:token",
		// TokenLookup: "cookie:token",

		// TokenHeadName is a string in the header. Default value is "Bearer"
		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,
	})

	if err != nil {
		log.Fatal("JWT Error:" + err.Error())
	}

	utils.jwt = jwtImpl

	return utils
}