kratos-utils/common/mwt.go

package common

import (
	"time"

	"github.com/vmihailenco/msgpack/v5"

	"github.com/dcsunny/mwt"
)

// secret:加密秘钥
// expire:过期时间,单位为秒
type MWT struct {
	secret []byte
	expire int64
}

type MWTClaims struct {
	AccountId string             `msgpack:"id"`
	Claims    msgpack.RawMessage `msgpack:"claims,omitempty"`
	mwt.StandardClaims
}

var (
	DefaultMWT *MWT
)

const (
	mwtHeader = "gqN0eXCjTVdUo2FsZ6VIUzI1Ng"
)

func InitMWT(secret string, expire int64) {
	DefaultMWT = NewMWT(secret, expire)
}

func NewMWT(secret string, expire int64) *MWT {
	var j = new(MWT)
	j.secret = []byte(secret)
	j.expire = expire
	return j
}

func (j *MWT) NewToken(accountID string, claimMap msgpack.RawMessage) (token string, tokenExpire int64, err error) {
	expireAt := time.Now().Add(time.Duration(j.expire) * time.Second).Unix()
	claims := MWTClaims{
		AccountId: accountID,
		Claims:    claimMap,
		StandardClaims: mwt.StandardClaims{
			ExpiresAt: expireAt,
		},
	}
	t := mwt.NewWithClaims(mwt.SigningMethodHS256, claims)
	token, err = t.SignedString(j.secret)
	if err != nil {
		return
	}
	token = token[27:]
	tokenExpire = claims.ExpiresAt
	return
}

func (j *MWT) Parse(text string) (userId string, expiresAt int64, claimMap msgpack.RawMessage, err error) {
	text = mwtHeader + "." + text
	token, err := mwt.ParseWithClaims(text, &MWTClaims{}, func(token *mwt.Token) (interface{}, error) {
		return j.secret, nil
	})
	if err != nil {
		if ve, ok := err.(*mwt.ValidationError); ok {
			if ve.Errors == mwt.ValidationErrorExpired {
				err = ErrTokenExpired
			}
		}
		return
	}
	if claims, ok := token.Claims.(*MWTClaims); ok {
		userId = claims.AccountId
		expiresAt = claims.StandardClaims.ExpiresAt
		claimMap = claims.Claims
		if token.Valid {
			return claims.AccountId, claims.StandardClaims.ExpiresAt, claims.Claims, nil
		}
	}
	err = ErrTokenInvalid
	return
}