gitee.com/h79/goutils@v1.22.10/auth/jwt/token.go

package jwt

import (
	"encoding/base64"
	"errors"
	"fmt"
	"gitee.com/h79/goutils/auth/token"
	"gitee.com/h79/goutils/common/result"
	"gitee.com/h79/goutils/common/timer"
	"github.com/golang-jwt/jwt/v5"
	"time"
)

type UpdateJWTClaims func(tk token.Token) jwt.Claims

// 保证 Token struct implement token.Token
var _ token.Token = (*Token)(nil)

type Token struct {
	Uid     string       `json:"uid"`
	Token   string       `json:"token"`
	Session string       `json:"-"` //for server and server
	Extend  string       `json:"-"`
	Source  token.Source `json:"-"`
	expire  timer.Expire
	ujc     UpdateJWTClaims
	key     token.Key
	custom  interface{}
	base64  bool
}

func NewToken(key token.Key, ujc UpdateJWTClaims, base64 bool) *Token {
	return &Token{
		ujc:    ujc,
		key:    key,
		base64: base64,
	}
}

func (tk *Token) WithVar(key string, value interface{}) token.Token {
	switch key {
	case token.CSource:
		tk.Source = value.(token.Source)
	case token.CSession:
		tk.Session = value.(string)
	case token.CExtend:
		tk.Extend = value.(string)
	case token.CUid:
		tk.Uid = value.(string)
	case token.CCustom:
		tk.custom = value
	case token.CIssuedAt:
		tk.expire.StartIn = value.(int64)
	case token.CExpiresAt:
		tk.expire.ExpireIn = value.(int64)
	}
	return tk
}

func (tk *Token) GetVar(key string) interface{} {
	switch key {
	case token.CSource:
		return tk.Source
	case token.CSession:
		return tk.Session
	case token.CUid:
		return tk.Uid
	case token.CBase64:
		return tk.base64
	case token.CCustom:
		return tk.custom
	case token.CExtend:
		return tk.Extend
	case token.CIssuedAt:
		return tk.expire.StartIn
	case token.CExpiresAt:
		return tk.expire.ExpireIn
	}
	return ""
}

func (tk *Token) GetToken() string {
	return tk.Token
}

// ExpireIn token.Token interface
func (tk *Token) ExpireIn() int64 {
	return tk.key.ExpireIn()
}

// Update token.Token interface
func (tk *Token) Update() (string, error) {
	jwtToken := jwt.NewWithClaims(jwt.SigningMethodHS256, tk.jwtClaims())
	tStr, err := jwtToken.SignedString(tk.key.Key())
	if err != nil {
		return "", err
	}
	if tk.base64 {
		tk.Token = base64.StdEncoding.EncodeToString([]byte(tStr))
	} else {
		tk.Token = tStr
	}
	return tk.Token, nil
}

func (tk *Token) jwtClaims() jwt.Claims {
	if tk.ujc != nil {
		return tk.ujc(tk)
	}
	var ex = tk.expire
	if tk.expire.StartIn <= 0 {
		ex = timer.NowExpireWithSecond(tk.ExpireIn())
	}
	return &jwtClaims{
		Subject:   tk.Uid,
		IssuedAt:  jwt.NewNumericDate(time.Unix(ex.StartIn, 0)),
		ExpiresAt: jwt.NewNumericDate(time.Unix(ex.ExpireIn, 0)),
		Source:    tk.Source,
		Session:   tk.Session,
		Extend:    tk.Extend,
	}
}

var defJWTFactory jwtFactory

type jwtFactory struct {
}

func (f *jwtFactory) Create(key token.Key, opts ...token.Option) (token.Token, error) {
	return NewToken(key, nil, token.IsBase64Enabled(opts...)), nil
}

func (f *jwtFactory) Decode(tk string, opts ...token.Option) (token.Token, error) {
	if token.IsBase64Enabled(opts...) {
		return DecodeToken(tk)
	}
	return DecodeTokenNoBase64(tk)
}

func (f *jwtFactory) Check(tk string, key token.Key, opts ...token.Option) (token.Token, error) {
	if token.IsBase64Enabled(opts...) {
		return VerifyToken(tk, key)
	}
	return VerifyTokenNoBase64(tk, key)
}

// DecodeToken 不需要认证，解包里面的信息
func DecodeToken(tk string) (token.Token, error) {
	bytes, er := base64.StdEncoding.DecodeString(tk)
	if er != nil {
		return nil, er
	}
	return DecodeTokenNoBase64(string(bytes))
}

func DecodeTokenNoBase64(tk string) (token.Token, error) {

	// for 读取信息,不验证
	tc := jwtClaims{}
	jtk, err := jwt.ParseWithClaims(tk, &tc, nil)
	if jtk != nil {
		return tc.To(tk), nil
	} else if errors.Is(err, jwt.ErrTokenMalformed) {
		fmt.Println("That's not even a token")
	} else if errors.Is(err, jwt.ErrTokenSignatureInvalid) {
		// Invalid signature
		fmt.Println("Invalid signature")
	} else if errors.Is(err, jwt.ErrTokenExpired) || errors.Is(err, jwt.ErrTokenNotValidYet) {
		// Token is either expired or not active yet
		fmt.Println("Timing is everything")
	} else {
		fmt.Println("Couldn't handle this token:", err)
	}
	return nil, err
}

// VerifyToken 需要认证，解包里面的信息
func VerifyToken(tk string, key token.Key) (token.Token, error) {
	bytes, er := base64.StdEncoding.DecodeString(tk)
	if er != nil {
		return nil, er
	}

	return VerifyTokenNoBase64(string(bytes), key)
}

func VerifyTokenNoBase64(tk string, key token.Key) (token.Token, error) {

	tc := jwtClaims{}
	jtk, err := jwt.ParseWithClaims(tk, &tc, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return key.Key(), nil
	})
	if err == nil && jtk.Valid {
		tok := tc.To(tk)
		tok.key = key
		return tok, nil
	}
	code := result.ErrToken
	if errors.Is(err, jwt.ErrTokenMalformed) {
		code |= result.TokenMalformed
	} else if errors.Is(err, jwt.ErrTokenUnverifiable) {
		code |= result.TokenUnVerify
	} else if errors.Is(err, jwt.ErrTokenSignatureInvalid) {
		code |= result.TokenSigned
	} else if errors.Is(err, jwt.ErrTokenExpired) || errors.Is(err, jwt.ErrTokenNotValidYet) {
		// Token is either expired or not active yet
		code |= result.TokenExpired
	} else {
		code |= result.TokenOther
	}
	return nil, result.ErrCode(code).WithError(err)
}

type jwtClaims struct {
	Source    token.Source     `json:"source"`
	Session   string           `json:"session"`
	Extend    string           `json:"extend"`
	Subject   string           `json:"sub,omitempty"`
	ExpiresAt *jwt.NumericDate `json:"exp,omitempty"`
	IssuedAt  *jwt.NumericDate `json:"iat,omitempty"`
}

// GetExpirationTime implements the Claims interface.
func (tc *jwtClaims) GetExpirationTime() (*jwt.NumericDate, error) {
	return tc.ExpiresAt, nil
}

// GetNotBefore implements the Claims interface.
func (tc *jwtClaims) GetNotBefore() (*jwt.NumericDate, error) {
	return nil, nil
}

// GetIssuedAt implements the Claims interface.
func (tc *jwtClaims) GetIssuedAt() (*jwt.NumericDate, error) {
	return tc.IssuedAt, nil
}

// GetAudience implements the Claims interface.
func (tc *jwtClaims) GetAudience() (jwt.ClaimStrings, error) {
	return jwt.ClaimStrings{}, nil
}

// GetIssuer implements the Claims interface.
func (tc *jwtClaims) GetIssuer() (string, error) {
	return "", nil
}

// GetSubject implements the Claims interface.
func (tc *jwtClaims) GetSubject() (string, error) {
	return tc.Subject, nil
}

func (tc *jwtClaims) To(tk string) *Token {
	return &Token{
		Session: tc.Session,
		Uid:     tc.Subject,
		Token:   tk,
		Extend:  tc.Extend,
		Source:  tc.Source,
		expire:  timer.Expire{StartIn: tc.IssuedAt.Unix(), ExpireIn: tc.ExpiresAt.Unix()},
	}
}