micro_libary_v4/token/token.go

package token

import (
	"context"
	"dashoo.cn/micro_libary/micro"
	"fmt"
	"github.com/gogf/gf/v2/crypto/gaes"
	"github.com/gogf/gf/v2/crypto/gmd5"
	"github.com/gogf/gf/v2/encoding/gbase64"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gctx"
	"github.com/gogf/gf/v2/os/gtime"
	"github.com/gogf/gf/v2/text/gstr"
	"github.com/gogf/gf/v2/util/gconv"
	"github.com/gogf/gf/v2/util/grand"
	"github.com/smallnest/rpcx/protocol"
	"strings"
)

// RpcxToken 结构体
type RpcxToken struct {
	// 缓存模式 1 gcache 2 gredis 默认1
	CacheMode int64
	// 缓存key
	CacheKey string
	// 超时时间 默认10天（毫秒）
	Timeout int64
	// 缓存刷新时间 默认为超时时间的一半（毫秒）
	MaxRefresh int64
	// Token分隔符
	TokenDelimiter string
	// Token加密key
	EncryptKey []byte
	// 认证失败中文提示
	AuthFailMsg string
	// 是否支持多端登录，默认false
	MultiLogin bool
	// 中间件类型 1 GroupMiddleware 2 BindMiddleware  3 GlobalMiddleware
	MiddlewareType int64

	// 拦截地址
	AuthPaths []string
	// 拦截排除地址
	AuthExcludePaths []string
	// 认证验证方法 return true 继续执行，否则结束执行
	AuthBeforeFunc func(req *protocol.Message) bool
	// 认证返回方法
	AuthAfterFunc func(req *protocol.Message, resp micro.Response)
}

func NewRpcxToken() *RpcxToken {
	m := new(RpcxToken)
	ctx := gctx.New()
	cacheMode, err := g.Config().Get(ctx, "token.cache-mode")
	if err != nil || cacheMode.Int64() == 0 {
		m.CacheMode = CacheModeCache
	} else {
		m.CacheMode = cacheMode.Int64()
	}

	cacheKey, err := g.Config().Get(ctx, "token.cache-key")
	if err != nil || cacheKey.String() == "" {
		m.CacheKey = DefaultCacheKey
	} else {
		m.CacheKey = cacheKey.String()
	}

	timeout, err := g.Config().Get(ctx, "token.timeout")
	if err != nil || timeout.Int64() == 0 {
		m.Timeout = DefaultTimeout
	} else {
		m.Timeout = timeout.Int64()
	}

	maxRefresh, err := g.Config().Get(ctx, "token.max-refresh")
	if err != nil || maxRefresh.Int64() == 0 {
		m.MaxRefresh = m.Timeout / 2
	} else {
		m.MaxRefresh = maxRefresh.Int64()
	}

	tokenDelimiter, err := g.Config().Get(ctx, "token.token-delimiter")
	if err != nil || tokenDelimiter.String() == "" {
		m.TokenDelimiter = DefaultTokenDelimiter
	} else {
		m.TokenDelimiter = tokenDelimiter.String()
	}

	encryptKey, err := g.Config().Get(ctx, "token.encrypt-key")
	if err != nil || len(encryptKey.String()) == 0 {
		m.EncryptKey = []byte(DefaultEncryptKey)
	} else {
		m.EncryptKey = []byte(encryptKey.String())
	}
	authFailMsg, err := g.Config().Get(ctx, "token.auth-fail-msg")
	if err != nil || authFailMsg.String() == "" {
		m.AuthFailMsg = DefaultAuthFailMsg
	} else {
		m.AuthFailMsg = authFailMsg.String()
	}

	multiLogin, err := g.Config().Get(ctx, "token.multi-login")
	if err != nil || multiLogin.String() == "" {
		m.MultiLogin = false
	} else {
		m.MultiLogin = authFailMsg.Bool()
	}

	return m
}

// GenToken 生成Token
func (m *RpcxToken) GenToken(ctx context.Context, userKey string, data interface{}) *micro.Response {
	token := m.EncryptToken(ctx, userKey, "")
	if !token.IsSuccess() {
		return token
	}

	cacheKey := m.CacheKey + userKey
	userCache := g.Map{
		KeyUserKey:     userKey,
		KeyUuid:        token.Get(KeyUuid).String(),
		KeyData:        data,
		KeyCreateTime:  gtime.Now().TimestampMilli(),
		KeyRefreshTime: gtime.Now().TimestampMilli() + gconv.Int64(m.MaxRefresh),
	}

	cacheResp := m.setCache(ctx, cacheKey, userCache)
	if !cacheResp.IsSuccess() {
		return cacheResp
	}

	return token
}

// GetTokenData 通过token获取对象
func (m *RpcxToken) GetTokenData(ctx context.Context, req *protocol.Message) *micro.Response {
	respData := m.GetRequestToken(ctx, req)
	if respData.IsSuccess() {
		// 验证token
		respData = m.validToken(ctx, respData.String())
	}

	return respData
}

// getRequestToken 返回请求Token
func (m *RpcxToken) GetRequestToken(ctx context.Context, req *protocol.Message) *micro.Response {
	authHeader := req.Metadata["Authorization"]
	if authHeader != "" {
		parts := strings.SplitN(authHeader, " ", 2)
		if !(len(parts) == 2 && parts[0] == "Bearer") {
			g.Log().Warning(ctx, msgLog(MsgErrAuthHeader, authHeader))
			return micro.Unauthorized(fmt.Sprintf(MsgErrAuthHeader, authHeader), "")
		} else if parts[1] == "" {
			g.Log().Warning(ctx, msgLog(MsgErrAuthHeader, authHeader))
			return micro.Unauthorized(fmt.Sprintf(MsgErrAuthHeader, authHeader), "")
		}

		return micro.Success(parts[1])
	}

	authHeader = req.Metadata[KeyToken]
	if authHeader == "" {
		return micro.Unauthorized(MsgErrTokenEmpty, "")
	}
	return micro.Success(authHeader)

}

// validToken 验证Token
func (m *RpcxToken) validToken(ctx context.Context, token string) *micro.Response {
	if token == "" {
		return micro.Unauthorized(MsgErrTokenEmpty, "")
	}

	decryptToken := m.DecryptToken(ctx, token)
	if !decryptToken.IsSuccess() {
		return decryptToken
	}

	userKey := decryptToken.Get(KeyUserKey)
	uuid := decryptToken.Get(KeyUuid)

	userCacheResp := m.GetToken(ctx, userKey.String())
	if !userCacheResp.IsSuccess() {
		return userCacheResp
	}

	if uuid.String() != userCacheResp.Get(KeyUuid).String() {
		g.Log().Debug(ctx, msgLog(MsgErrAuthUuid)+", decryptToken:"+decryptToken.Json()+" cacheValue:"+gconv.String(userCacheResp.Data))
		return micro.Unauthorized(MsgErrAuthUuid, "")
	}

	return userCacheResp
}

// GetToken 通过userKey获取Token
func (m *RpcxToken) GetToken(ctx context.Context, userKey string) *micro.Response {
	cacheKey := m.CacheKey + userKey

	userCacheResp := m.getCache(ctx, cacheKey)
	if !userCacheResp.IsSuccess() {
		return userCacheResp
	}
	userCache := gconv.Map(userCacheResp.Data)

	nowTime := gtime.Now().TimestampMilli()
	refreshTime := userCache[KeyRefreshTime]

	// 需要进行缓存超时时间刷新
	if gconv.Int64(refreshTime) == 0 || nowTime > gconv.Int64(refreshTime) {
		userCache[KeyCreateTime] = gtime.Now().TimestampMilli()
		userCache[KeyRefreshTime] = gtime.Now().TimestampMilli() + gconv.Int64(m.MaxRefresh)
		return m.setCache(ctx, cacheKey, userCache)
	}

	return micro.Success(userCache)
}

// RemoveToken 删除Token
func (m *RpcxToken) RemoveToken(ctx context.Context, token string) *micro.Response {
	decryptToken := m.DecryptToken(ctx, token)
	if !decryptToken.IsSuccess() {
		return decryptToken
	}

	cacheKey := m.CacheKey + decryptToken.Get(KeyUserKey).String()
	return m.removeCache(ctx, cacheKey)
}

// EncryptToken token加密方法
func (m *RpcxToken) EncryptToken(ctx context.Context, userKey string, uuid string) *micro.Response {
	if userKey == "" {
		return micro.Fail(MsgErrUserKeyEmpty)
	}

	if uuid == "" {
		// 重新生成uuid
		newUuid, err := gmd5.Encrypt(grand.Letters(10))
		if err != nil {
			g.Log().Error(ctx, msgLog(MsgErrAuthUuid), err)
			return micro.Error(MsgErrAuthUuid)
		}
		uuid = newUuid
	}

	tokenStr := userKey + m.TokenDelimiter + uuid

	token, err := gaes.Encrypt([]byte(tokenStr), m.EncryptKey)
	if err != nil {
		g.Log().Error(ctx, msgLog(MsgErrTokenEncrypt), tokenStr, err)
		return micro.Error(MsgErrTokenEncrypt)
	}

	return micro.Success(g.Map{
		//KeyUserKey: userKey,
		KeyUuid:  uuid,
		KeyToken: gbase64.EncodeToString(token),
	})
}

// DecryptToken token解密方法
func (m *RpcxToken) DecryptToken(ctx context.Context, token string) *micro.Response {
	if token == "" {
		return micro.Fail(MsgErrTokenEmpty)
	}

	token64, err := gbase64.Decode([]byte(token))
	if err != nil {
		g.Log().Error(ctx, msgLog(MsgErrTokenDecode), token, err)
		return micro.Error(MsgErrTokenDecode)
	}
	decryptToken, err2 := gaes.Decrypt(token64, m.EncryptKey)
	if err2 != nil {
		g.Log().Error(ctx, msgLog(MsgErrTokenEncrypt), token, err2)
		return micro.Error(MsgErrTokenEncrypt)
	}
	tokenArray := gstr.Split(string(decryptToken), m.TokenDelimiter)
	if len(tokenArray) < 2 {
		g.Log().Error(ctx, msgLog(MsgErrTokenLen), token)
		return micro.Error(MsgErrTokenLen)
	}

	return micro.Success(g.Map{
		KeyUserKey: tokenArray[0],
		KeyUuid:    tokenArray[1],
	})
}

// String token解密方法
func (m *RpcxToken) String() string {
	return gconv.String(g.Map{
		// 缓存模式 1 gcache 2 gredis 默认1
		"CacheMode":        m.CacheMode,
		"CacheKey":         m.CacheKey,
		"Timeout":          m.Timeout,
		"TokenDelimiter":   m.TokenDelimiter,
		"AuthFailMsg":      m.AuthFailMsg,
		"MultiLogin":       m.MultiLogin,
		"MiddlewareType":   m.MiddlewareType,
		"AuthPaths":        gconv.String(m.AuthPaths),
		"AuthExcludePaths": gconv.String(m.AuthExcludePaths),
	})
}