From d0c2c27071c1429862f5e8007e15f17620fc52b8 Mon Sep 17 00:00:00 2001
From: JustSong <quanpengsong@gmail.com>
Date: Tue, 22 Nov 2022 15:11:02 +0800
Subject: [PATCH] fix: check if the token is shared before removing

---
 channel/token-store.go | 18 ++++++++++++------
 model/user.go          | 11 +++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/channel/token-store.go b/channel/token-store.go
index c5b2c84..98b35a3 100644
--- a/channel/token-store.go
+++ b/channel/token-store.go
@@ -92,7 +92,6 @@ func TokenStoreRemoveItem(item TokenStoreItem) {
 }
 
 func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
-	// TODO: check if this token is shared!
 	if cleanUser.WeChatTestAccountId == originUser.WeChatTestAccountId {
 		cleanUser.WeChatTestAccountId = ""
 	}
@@ -111,7 +110,9 @@ func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
 		if cleanUser.WeChatTestAccountSecret != "" {
 			newWeChatTestAccountTokenStoreItem.AppSecret = cleanUser.WeChatTestAccountSecret
 		}
-		TokenStoreRemoveItem(&oldWeChatTestAccountTokenStoreItem)
+		if !model.IsWeChatTestAccountTokenShared(&oldWeChatTestAccountTokenStoreItem) {
+			TokenStoreRemoveItem(&oldWeChatTestAccountTokenStoreItem)
+		}
 		TokenStoreAddItem(&newWeChatTestAccountTokenStoreItem)
 	}
 	if cleanUser.WeChatCorpAccountId == originUser.WeChatCorpAccountId {
@@ -139,25 +140,30 @@ func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
 		if cleanUser.WeChatCorpAccountAgentId != "" {
 			newWeChatCorpAccountTokenStoreItem.AgentId = cleanUser.WeChatCorpAccountAgentId
 		}
-		TokenStoreRemoveItem(&oldWeChatCorpAccountTokenStoreItem)
+		if !model.IsWeChatCorpAccountTokenShared(&oldWeChatCorpAccountTokenStoreItem) {
+			TokenStoreRemoveItem(&oldWeChatCorpAccountTokenStoreItem)
+		}
 		TokenStoreAddItem(&newWeChatCorpAccountTokenStoreItem)
 	}
 }
 
 // TokenStoreRemoveUser user must be filled
 func TokenStoreRemoveUser(user *model.User) {
-	// TODO: check if this token is shared!
 	testAccountTokenStoreItem := WeChatTestAccountTokenStoreItem{
 		AppID:     user.WeChatTestAccountId,
 		AppSecret: user.WeChatTestAccountSecret,
 	}
-	TokenStoreRemoveItem(&testAccountTokenStoreItem)
+	if !model.IsWeChatTestAccountTokenShared(&testAccountTokenStoreItem) {
+		TokenStoreRemoveItem(&testAccountTokenStoreItem)
+	}
 	corpAccountTokenStoreItem := WeChatCorpAccountTokenStoreItem{
 		CorpId:     user.WeChatCorpAccountId,
 		CorpSecret: user.WeChatCorpAccountSecret,
 		AgentId:    user.WeChatCorpAccountAgentId,
 	}
-	TokenStoreRemoveItem(&corpAccountTokenStoreItem)
+	if !model.IsWeChatCorpAccountTokenShared(&corpAccountTokenStoreItem) {
+		TokenStoreRemoveItem(&corpAccountTokenStoreItem)
+	}
 }
 
 func TokenStoreGetToken(key string) string {
diff --git a/model/user.go b/model/user.go
index c84136a..91be1c4 100644
--- a/model/user.go
+++ b/model/user.go
@@ -2,6 +2,7 @@ package model
 
 import (
 	"errors"
+	"message-pusher/channel"
 	"message-pusher/common"
 	"strings"
 )
@@ -168,6 +169,16 @@ func IsUsernameAlreadyTaken(username string) bool {
 	return DB.Where("username = ?", username).Find(&User{}).RowsAffected == 1
 }
 
+func IsWeChatTestAccountTokenShared(item *channel.WeChatTestAccountTokenStoreItem) bool {
+	return DB.Where("wechat_test_account_id = ? and wechat_test_account_secret = ?",
+		item.AppID, item.AppSecret).Find(&User{}).RowsAffected != 1
+}
+
+func IsWeChatCorpAccountTokenShared(item *channel.WeChatCorpAccountTokenStoreItem) bool {
+	return DB.Where("wechat_corp_account_id = ? and wechat_corp_account_secret = ? and wechat_corp_account_agent_id = ?",
+		item.CorpId, item.CorpSecret, item.AgentId).Find(&User{}).RowsAffected != 1
+}
+
 func ResetUserPasswordByEmail(email string, password string) error {
 	hashedPassword, err := common.Password2Hash(password)
 	if err != nil {