From 6d4edc1f5b6d01a1a8ca622e7e58a82ad68a8bf6 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Sat, 30 Nov 2024 23:32:42 +0800
Subject: [PATCH] fix: realtime

---
 middleware/auth.go         | 1 +
 relay/common/relay_info.go | 3 +++
 service/quota.go           | 2 +-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/middleware/auth.go b/middleware/auth.go
index 53c7079d..cb55cac9 100644
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -212,6 +212,7 @@ func TokenAuth() func(c *gin.Context) {
 		}
 		c.Set("id", token.UserId)
 		c.Set("token_id", token.Id)
+		c.Set("token_key", token.Key)
 		c.Set("token_name", token.Name)
 		c.Set("token_unlimited_quota", token.UnlimitedQuota)
 		if !token.UnlimitedQuota {
diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go
index 76ae0367..ae95cd51 100644
--- a/relay/common/relay_info.go
+++ b/relay/common/relay_info.go
@@ -14,6 +14,7 @@ type RelayInfo struct {
 	ChannelType          int
 	ChannelId            int
 	TokenId              int
+	TokenKey             string
 	UserId               int
 	Group                string
 	TokenUnlimited       bool
@@ -58,6 +59,7 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 	channelId := c.GetInt("channel_id")
 
 	tokenId := c.GetInt("token_id")
+	tokenKey := c.GetString("token_key")
 	userId := c.GetInt("id")
 	group := c.GetString("group")
 	tokenUnlimited := c.GetBool("token_unlimited_quota")
@@ -73,6 +75,7 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 		ChannelType:       channelType,
 		ChannelId:         channelId,
 		TokenId:           tokenId,
+		TokenKey:          tokenKey,
 		UserId:            userId,
 		Group:             group,
 		TokenUnlimited:    tokenUnlimited,
diff --git a/service/quota.go b/service/quota.go
index 0ca91457..70913306 100644
--- a/service/quota.go
+++ b/service/quota.go
@@ -22,7 +22,7 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
 		return err
 	}
 
-	token, err := model.CacheGetTokenByKey(strings.TrimLeft(relayInfo.ApiKey, "sk-"))
+	token, err := model.CacheGetTokenByKey(strings.TrimLeft(relayInfo.TokenKey, "sk-"))
 	if err != nil {
 		return err
 	}