user_token.go 1.6 KB
package contextdata

import (
	"context"
	"encoding/json"
	"github.com/golang-jwt/jwt/v4"
	"github.com/zeromicro/go-zero/core/logx"
	"time"
)

var (
	CtxKeyJwtUserId    = "userId"
	CtxKeyJwtCompanyId = "companyId"
)

func GetInt64FromCtx(ctx context.Context, key string) int64 {
	var uid int64
	if jsonUid, ok := ctx.Value(key).(json.Number); ok {
		if int64Uid, err := jsonUid.Int64(); err == nil {
			uid = int64Uid
		} else {
			logx.WithContext(ctx).Errorf("GetUidFromCtx err : %+v", err)
		}
	}
	return uid
}

func getStringFromCtx(ctx context.Context, key string) string {
	var uid string
	if jsonUid, ok := ctx.Value(key).(string); ok {
		return jsonUid
	}
	return uid
}

func getArrayInt64FromCtx(ctx context.Context, key string) []int64 {
	values := ctx.Value(key)
	var ids = make([]int64, 0)
	if values == nil {
		return ids
	}
	if list, ok := values.([]interface{}); ok {
		for _, item := range list {
			if jsonId, ok := item.(json.Number); ok {
				id, _ := jsonId.Int64()
				ids = append(ids, id)
			}
		}
	}
	return ids
}

func GetUserTokenFromCtx(ctx context.Context) UserToken {
	return UserToken{
		UserId:    GetInt64FromCtx(ctx, CtxKeyJwtUserId),
		CompanyId: GetInt64FromCtx(ctx, CtxKeyJwtCompanyId),
	}
}

type UserToken struct {
	UserId    int64 `json:"userId"`
	CompanyId int64 `json:"companyId"`
}

func (tk UserToken) GenerateToken(secret string, expire int64) (string, error) {
	claims := make(jwt.MapClaims)
	claims["exp"] = time.Now().Unix() + expire
	claims["iat"] = time.Now().Unix()
	claims["UserId"] = tk.UserId
	token := jwt.New(jwt.SigningMethodHS256)
	token.Claims = claims

	return token.SignedString([]byte(secret))
}