update

cmd/haobase/message/ws/client.go

@@ -8,10 +8,13 @@ import (
 	"bytes"
 	"encoding/json"
 	"net/http"
+	"strings"
 	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
+	"github.com/yzimhao/trading_engine/types"
+	"github.com/yzimhao/trading_engine/types/token"
 	"github.com/yzimhao/trading_engine/utils/app"
 )
 
@@ -144,7 +147,22 @@ func (c *Client) handleRecvData(body []byte) {
 	}
 
 	for _, attr := range msg.Subsc {
-		c.setAttr(attr)
+		if strings.HasPrefix(attr, "_") {
+			//带有_标记的tag只能是内部程序设置的，不能通过前端发送过来指定
+			continue
+		}
+		if strings.HasPrefix(attr, "token.") {
+			a := strings.Split(attr, ".")
+			_token := a[1]
+			user_id := token.Get(_token)
+			if user_id != "" {
+				c.setAttr(types.MsgUser.Format(map[string]string{
+					"user_id": user_id,
+				}))
+			}
+		} else {
+			c.setAttr(attr)
+		}
 	}
 
 	app.Logger.Debugf("[wss] recv: %v", msg)

cmd/haobase/www/internal_api/token.go

@@ -1,13 +1,10 @@
 package internal_api
 
 import (
-	"fmt"
-
 	"github.com/gin-gonic/gin"
-	"github.com/gomodule/redigo/redis"
 	"github.com/gookit/goutil/arrutil"
+	"github.com/yzimhao/trading_engine/types/token"
 	"github.com/yzimhao/trading_engine/utils"
-	"github.com/yzimhao/trading_engine/utils/app"
 	"github.com/yzimhao/trading_engine/utils/app/config"
 )
 
@@ -23,33 +20,9 @@ func SetToken(ctx *gin.Context) {
 		utils.ResponseFailJson(ctx, err.Error())
 		return
 	}
-	UpdateRedisToken(req)
-	utils.ResponseOkJson(ctx, "")
-}
-
-func UpdateRedisToken(req req_settoken_args) {
-	rdc := app.RedisPool().Get()
-	defer rdc.Close()
-
-	topic := tokenRedisTopic(req.Token)
-	rdc.Do("set", topic, req.UserId)
-	rdc.Do("expire", topic, req.Ttl)
-}
 
-func GetUserIdFromToken(original_token string) string {
-	rdc := app.RedisPool().Get()
-	defer rdc.Close()
-
-	topic := tokenRedisTopic(original_token)
-	user_id, err := redis.String(rdc.Do("get", topic))
-	if err != nil {
-		app.Logger.Errorf("从redis获取token信息出错 %s", err.Error())
-	}
-	return user_id
-}
-
-func tokenRedisTopic(token string) string {
-	return fmt.Sprintf("user.token.%s", utils.Hash256(token))
+	token.Set(req.Token, req.UserId, req.Ttl)
+	utils.ResponseOkJson(ctx, "")
 }
 
 func Authentication() gin.HandlerFunc {

cmd/haobase/www/middle/base.go

@@ -5,7 +5,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/gookit/goutil/arrutil"
-	"github.com/yzimhao/trading_engine/cmd/haobase/www/internal_api"
+	"github.com/yzimhao/trading_engine/types/token"
 	"github.com/yzimhao/trading_engine/utils"
 	"github.com/yzimhao/trading_engine/utils/app"
 	"github.com/yzimhao/trading_engine/utils/app/config"
@@ -18,9 +18,9 @@ var (
 func CheckLogin() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		user_id := ""
-		token := c.GetHeader("Token")
+		_token := c.GetHeader("Token")
 		if config.App.Main.Mode == config.ModeDemo {
-			user_id = token
+			user_id = _token
 			if user_id == "" {
 				user_id = c.Query("user_id")
 			}
@@ -43,7 +43,7 @@ func CheckLogin() gin.HandlerFunc {
 			}
 		} else {
 			//从redis的token中获取登陆用户ID
-			user_id = internal_api.GetUserIdFromToken(token)
+			user_id = token.Get(_token)
 		}
 
 		if user_id == "" {

types/token/token.go

@@ -0,0 +1,34 @@
+package token
+
+import (
+	"fmt"
+
+	"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) {
+	rdc := app.RedisPool().Get()
+	defer rdc.Close()
+
+	topic := tokenRedisTopic(token)
+	rdc.Do("set", topic, user_id)
+	rdc.Do("expire", topic, ttl)
+}
+
+func Get(original_token string) string {
+	rdc := app.RedisPool().Get()
+	defer rdc.Close()
+
+	topic := tokenRedisTopic(original_token)
+	user_id, err := redis.String(rdc.Do("get", topic))
+	if err != nil {
+		app.Logger.Errorf("从redis获取token信息出错 %s", err.Error())
+	}
+	return user_id
+}
+
+func tokenRedisTopic(token string) string {
+	return fmt.Sprintf("user.token.%s", utils.Hash256(token))
+}