From 47933c138f893e98b191d2f99b0bb04e126c4cd0 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Thu, 23 Nov 2023 16:58:30 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84token=E9=83=A8=E5=88=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/haobase/www/internal_api/token.go |  6 +++++-
 types/token/token.go                  | 17 ++++++++++++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/cmd/haobase/www/internal_api/token.go b/cmd/haobase/www/internal_api/token.go
index 5fb4b1f..2e89cd9 100644
--- a/cmd/haobase/www/internal_api/token.go
+++ b/cmd/haobase/www/internal_api/token.go
@@ -21,7 +21,11 @@ func SetToken(ctx *gin.Context) {
 		return
 	}
 
-	token.Set(req.Token, req.UserId, req.Ttl)
+	err := token.Set(req.Token, req.UserId, req.Ttl)
+	if err != nil {
+		utils.ResponseFailJson(ctx, err.Error())
+		return
+	}
 	utils.ResponseOkJson(ctx, "")
 }
 
diff --git a/types/token/token.go b/types/token/token.go
index 6b656af..521b003 100644
--- a/types/token/token.go
+++ b/types/token/token.go
@@ -2,19 +2,30 @@ package token
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/gomodule/redigo/redis"
 	"github.com/yzimhao/trading_engine/utils"
 	"github.com/yzimhao/trading_engine/utils/app"
 )
 
-func Set(token string, user_id string, ttl int) {
+func Set(token string, user_id string, ttl int) error {
 	rdc := app.RedisPool().Get()
 	defer rdc.Close()
 
+	if strings.Contains(token, ".") {
+		return fmt.Errorf("token contains `.`")
+	}
+
 	topic := tokenRedisTopic(token)
-	rdc.Do("set", topic, user_id)
-	rdc.Do("expire", topic, ttl)
+	if _, err := rdc.Do("set", topic, user_id); err != nil {
+		return err
+	}
+	if _, err := rdc.Do("expire", topic, ttl); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 func Get(original_token string) string {