Преглед изворни кода

feat: 更新令牌生成算法

[email protected] пре 1 година
родитељ
комит
b58b1dc0ec
3 измењених фајлова са 50 додато и 16 уклоњено
  1. 18 13
      common/utils.go
  2. 10 1
      controller/token.go
  3. 22 2
      controller/user.go

+ 18 - 13
common/utils.go

@@ -1,10 +1,12 @@
 package common
 
 import (
+	crand "crypto/rand"
 	"fmt"
 	"github.com/google/uuid"
 	"html/template"
 	"log"
+	"math/big"
 	"math/rand"
 	"net"
 	"os/exec"
@@ -145,21 +147,24 @@ func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
-func GenerateKey() string {
-	//rand.Seed(time.Now().UnixNano())
-	key := make([]byte, 48)
-	for i := 0; i < 16; i++ {
-		key[i] = keyChars[rand.Intn(len(keyChars))]
-	}
-	uuid_ := GetUUID()
-	for i := 0; i < 32; i++ {
-		c := uuid_[i]
-		if i%2 == 0 && c >= 'a' && c <= 'z' {
-			c = c - 'a' + 'A'
+func GenerateRandomKey(length int) (string, error) {
+	b := make([]byte, length)
+	maxI := big.NewInt(int64(len(keyChars)))
+
+	for i := range b {
+		n, err := crand.Int(crand.Reader, maxI)
+		if err != nil {
+			return "", err
 		}
-		key[i+16] = c
+		b[i] = keyChars[n.Int64()]
 	}
-	return string(key)
+
+	return string(b), nil
+}
+
+func GenerateKey() (string, error) {
+	//rand.Seed(time.Now().UnixNano())
+	return GenerateRandomKey(48)
 }
 
 func GetRandomInt(max int) int {

+ 10 - 1
controller/token.go

@@ -123,10 +123,19 @@ func AddToken(c *gin.Context) {
 		})
 		return
 	}
+	key, err := common.GenerateKey()
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "生成令牌失败",
+		})
+		common.SysError("failed to generate token key: " + err.Error())
+		return
+	}
 	cleanToken := model.Token{
 		UserId:             c.GetInt("id"),
 		Name:               token.Name,
-		Key:                common.GenerateKey(),
+		Key:                key,
 		CreatedTime:        common.GetTimestamp(),
 		AccessedTime:       common.GetTimestamp(),
 		ExpiredTime:        token.ExpiredTime,

+ 22 - 2
controller/user.go

@@ -200,11 +200,20 @@ func Register(c *gin.Context) {
 	}
 	// 生成默认令牌
 	if constant.GenerateDefaultToken {
+		key, err := common.GenerateKey()
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "生成默认令牌失败",
+			})
+			common.SysError("failed to generate token key: " + err.Error())
+			return
+		}
 		// 生成默认令牌
 		token := model.Token{
 			UserId:             insertedUser.Id, // 使用插入后的用户ID
 			Name:               cleanUser.Username + "的初始令牌",
-			Key:                common.GenerateKey(),
+			Key:                key,
 			CreatedTime:        common.GetTimestamp(),
 			AccessedTime:       common.GetTimestamp(),
 			ExpiredTime:        -1,     // 永不过期
@@ -311,7 +320,18 @@ func GenerateAccessToken(c *gin.Context) {
 		})
 		return
 	}
-	user.SetAccessToken(common.GetUUID())
+	// get rand int 28-32
+	randI := common.GetRandomInt(4)
+	key, err := common.GenerateRandomKey(29 + randI)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "生成失败",
+		})
+		common.SysError("failed to generate key: " + err.Error())
+		return
+	}
+	user.SetAccessToken(key)
 
 	if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 {
 		c.JSON(http.StatusOK, gin.H{