kit/core/token/token.go

package token

import (
	"fmt"
	"git.diulo.com/mogfee/kit/errors"
	"git.diulo.com/mogfee/kit/internal/xuuid"
	"github.com/golang-jwt/jwt/v5"
	"time"
)

// iss    : jwt签发者
// sub  : 主题名称
// aud  : 面向的用户，一般都是通过ip或者域名控制
// exp  : jwt的有效时间（过期时间），这个有效时间必须要大于签发时间，对于交互接口来说，建议是预设5秒
// nbf   : 在什么时候jwt开始生效（在此之前不可用）
// iat    : jwt的签发时间
// jti     : 唯一标识，主要用来回避被重复使用攻击

type (
	Options func(opt *tokenServer)

	tokenServer struct {
		prefix      string
		key         string
		expire      time.Duration
		generateKey func() string
		store       Store
	}
)

var (
	ErrTokenExpired = errors.Unauthorized("TOKEN_EXPIRED", "")
	ErrTokenError   = errors.Unauthorized("TOKEN_ERROR", "")
)

func WithPrefix(prefix string) Options {
	return func(opt *tokenServer) {
		opt.prefix = prefix
	}
}
func WithExpire(expire time.Duration) Options {
	return func(opt *tokenServer) {
		opt.expire = expire
	}
}
func WithGenerateKey(fn func() string) Options {
	return func(opt *tokenServer) {
		opt.generateKey = fn
	}
}
func WithStore(store Store) Options {
	return func(opt *tokenServer) {
		opt.store = store
	}
}
func NewTokenService(key string, opts ...Options) *tokenServer {
	srv := &tokenServer{
		key: key,
		generateKey: func() string {
			return xuuid.UUID()
		},
		store:  newStore(),
		prefix: "token:",
	}
	for _, opt := range opts {
		opt(srv)
	}
	return srv
}
func (t *tokenServer) getStoreKey(uniqId string) string {
	return fmt.Sprintf("%s%s", t.prefix, uniqId)
}
func (t *tokenServer) Generate(info any) (string, error) {
	token, uniqId, err := t.generate()
	if err != nil {
		return "", err
	}
	if err = t.store.Set(t.getStoreKey(uniqId), info, t.expire); err != nil {
		return "", err
	}
	return token, nil
}

func (t *tokenServer) generate() (token string, uniqId string, err error) {
	uniqId = t.generateKey()
	var tokenStr string
	ctime := time.Now().Unix()
	tokenStr, err = jwt.NewWithClaims(jwt.SigningMethodHS256, &jwt.MapClaims{
		//jwt签发者
		//"iss": "auth.diulo.com",
		//主题名称
		//"sub": "www.diulo.com",
		// 面向的用户，一般都是通过ip或者域名控制
		//"aud": "api.diulo.com",
		"aud": fmt.Sprint(uniqId),
		//jwt的有效时间（过期时间），这个有效时间必须要大于签发时间，对于交互接口来说，建议是预设5秒
		//"exp": ctime + 7200,
		//在什么时候jwt开始生效（在此之前不可用）
		"nbf": ctime,
		//jwt的签发时间
		"iat": ctime,
		//唯一标识，主要用来回避被重复使用攻击
		//"jti": fmt.Sprint(uniqId),
	}).SignedString([]byte(t.key))
	if err != nil {
		return
	}
	token, err = Encrypt(tokenStr, []byte(t.key), t.key)
	return
}

func (t *tokenServer) Refresh(token string) error {
	uniqId, info, err := t.parse(token)
	if err != nil {
		return err
	}
	return t.store.Set(uniqId, info, t.expire)
}
func (t *tokenServer) Parse(tokenStr string) (any, error) {
	_, info, err := t.parse(tokenStr)
	return info, err
}
func (t *tokenServer) parse(tokenStr string) (string, any, error) {
	var res any
	if tokenStr == "" {
		return "", res, nil
	}
	str, err := Decrypt(tokenStr, []byte(t.key), t.key)
	if err != nil {
		return "", res, ErrTokenError
	}
	token, err := jwt.Parse(str, func(token *jwt.Token) (interface{}, error) {
		return []byte(t.key), nil
	})
	if err != nil {
		if errors.Is(err, jwt.ErrTokenExpired) {
			return "", res, ErrTokenExpired
		}
		return "", res, ErrTokenError
	}

	if token.Valid {
		audience, err := token.Claims.GetAudience()
		if err != nil {
			return "", res, ErrTokenError
		}
		if len(audience) > 0 {
			if info, ok := t.store.Get(audience[0]); ok {
				return audience[0], info, nil
			}
		}
		return "", res, nil
	} else if errors.Is(err, jwt.ErrTokenExpired) || errors.Is(err, jwt.ErrTokenNotValidYet) {
		return "", res, ErrTokenExpired
	}
	return "", res, ErrTokenError
}