From 353ada134392cb8c15e86a0489fd2ca700f5f3bd Mon Sep 17 00:00:00 2001
From: Yun <yun@yun.ink>
Date: Sat, 23 Mar 2024 20:23:08 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E7=BB=93=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod       |  2 ++
 go.sum       |  2 ++
 jwtx.go      | 67 ++++++++++++++++++++++++++--------------------------
 jwtx_test.go | 28 ++++++++++++++++++++++
 4 files changed, 65 insertions(+), 34 deletions(-)
 create mode 100644 go.sum
 create mode 100644 jwtx_test.go

diff --git a/go.mod b/go.mod
index d284078..3709a42 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module code.yun.ink/pkg/jwtx
 
 go 1.20
+
+require github.com/dgrijalva/jwt-go v3.2.0+incompatible
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..6a8f140
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
+github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
diff --git a/jwtx.go b/jwtx.go
index dec8d39..d82f50d 100644
--- a/jwtx.go
+++ b/jwtx.go
@@ -1,6 +1,7 @@
 package jwtx
 
 import (
+	"encoding/json"
 	"errors"
 	"time"
 
@@ -8,80 +9,78 @@ import (
 )
 
 var (
-	TokenExpired        error = errors.New("token is expired")
-	TokenNotValidYet    error = errors.New("token not active yet")
-	TokenMalformed      error = errors.New("that's not even a token")
-	TokenInvalid        error = errors.New("couldn't handle this token")
-	TokenGenerateFailed error = errors.New("generate token failed!")
-
-	// defaultSignKey string = "hcVI@Nm4cjhjPVvCpL5^1%jY"
-	defaultSignKey = "XiaoXinPlus"
-
-	RedisUserKey string = "user:%d"
+	ErrTokenExpired        error = errors.New("token is expired")
+	ErrTokenNotValidYet    error = errors.New("token not active yet")
+	ErrTokenMalformed      error = errors.New("that's not even a token")
+	ErrTokenInvalid        error = errors.New("couldn't handle this token")
+	ErrTokenGenerateFailed error = errors.New("generate token failed!")
 )
 
 type jwtx struct {
-	SigningKey []byte
+	SignKey []byte
 }
 
 type Claims struct {
-	TenantId int64 `json:"tenant_id"`
-	ClientId int64 `json:"client_id"`
-	UserId   int64 `json:"user_id"`
-
-	Tag string `json:"tag"` // 额外标记 
-
+	UserData string `json:"user_data"`
 	jwt.StandardClaims
 }
 
-func NewJWT() *jwtx {
+func NewJWT(signKey string) *jwtx {
 	return &jwtx{
-		[]byte(defaultSignKey),
+		[]byte(signKey),
 	}
 }
 
 // 仅解析Token
-func (j *jwtx) Decode(tokenString string) (*Claims, error) {
+func (j *jwtx) Decode(tokenString string, userData interface{}) error {
 	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
-		return j.SigningKey, nil
+		return j.SignKey, nil
 	})
 	if err != nil {
 		if ve, ok := err.(*jwt.ValidationError); ok {
 			if ve.Errors&jwt.ValidationErrorMalformed == jwt.ValidationErrorMalformed {
-				return nil, TokenMalformed
+				return ErrTokenMalformed
 			} else if ve.Errors&jwt.ValidationErrorExpired == jwt.ValidationErrorExpired {
 				// Token is expired
-				return nil, TokenExpired
+				return ErrTokenExpired
 			} else if ve.Errors&jwt.ValidationErrorNotValidYet == jwt.ValidationErrorNotValidYet {
-				return nil, TokenNotValidYet
+				return ErrTokenNotValidYet
 			} else {
-				return nil, TokenInvalid
+				return ErrTokenInvalid
 			}
 		}
 	}
 	claims, ok := token.Claims.(*Claims)
 	if !ok || !token.Valid {
-		return nil, TokenInvalid
+		return ErrTokenInvalid
 	}
-	return claims, nil
+
+	return json.Unmarshal([]byte(claims.UserData), userData)
+
 }
 
 // 根据参数生成和设置 token
-func (j *jwtx) Encode(claims *Claims, expiresTime time.Time) (string, error) {
+func (j *jwtx) Encode(userData interface{}, expiresTime time.Time) (string, error) {
 
-	expSecond := expiresTime.Unix()
-	nowSecond := time.Now().Unix()
-
-	if expSecond-nowSecond < 0 {
+	if expiresTime.Before(time.Now()) {
 		return "", errors.New("时间无效")
 	}
 
+	db, err := json.Marshal(userData)
+	if err != nil {
+		return "", errors.New("json序列化失败")
+	}
+
+	claims := Claims{
+		UserData: string(db),
+	}
+
 	// 设置有效时间
 	claims.StandardClaims.ExpiresAt = expiresTime.Unix()
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
-	tokenString, err := token.SignedString(j.SigningKey)
+	tokenString, err := token.SignedString(j.SignKey)
 	if err != nil {
-		return "", TokenGenerateFailed
+		return "", ErrTokenGenerateFailed
 	}
 	return tokenString, nil
 }
diff --git a/jwtx_test.go b/jwtx_test.go
new file mode 100644
index 0000000..aa3d206
--- /dev/null
+++ b/jwtx_test.go
@@ -0,0 +1,28 @@
+package jwtx_test
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"code.yun.ink/pkg/jwtx"
+)
+
+func TestJwt(t *testing.T) {
+	u := UserData{
+		UserId:   1,
+		UserName: "test",
+	}
+	j := jwtx.NewJWT("test")
+	s, err := j.Encode(u, time.Now().AddDate(0, 0, 1))
+	fmt.Println(s, err)
+	// var u2 UserData
+	u2 := UserData{}
+	err = j.Decode(s, &u2)
+	fmt.Println(u2, err)
+}
+
+type UserData struct {
+	UserId   int64  `json:"user_id"`
+	UserName string `json:"user_name"`
+}