fix: fix token store not correctly registered (#30)

channel/token-store.go

@@ -11,6 +11,7 @@ type TokenStoreItem interface {
 	Key() string
 	Token() string
 	Refresh()
+	IsFilled() bool
 }
 
 type tokenStore struct {
@@ -78,7 +79,11 @@ func TokenStoreInit() {
 	}()
 }
 
+// TokenStoreAddItem It's okay to add an incomplete item.
 func TokenStoreAddItem(item TokenStoreItem) {
+	if !item.IsFilled() {
+		return
+	}
 	item.Refresh()
 	s.Mutex.RLock()
 	s.Map[item.Key()] = &item
@@ -91,18 +96,36 @@ func TokenStoreRemoveItem(item TokenStoreItem) {
 	s.Mutex.RUnlock()
 }
 
+func TokenStoreAddUser(user *model.User) {
+	testItem := WeChatTestAccountTokenStoreItem{
+		AppID:     user.WeChatTestAccountId,
+		AppSecret: user.WeChatTestAccountSecret,
+	}
+	TokenStoreAddItem(&testItem)
+	corpItem := WeChatCorpAccountTokenStoreItem{
+		CorpId:      user.WeChatCorpAccountId,
+		AgentSecret: user.WeChatCorpAccountAgentSecret,
+		AgentId:     user.WeChatCorpAccountAgentId,
+	}
+	TokenStoreAddItem(&corpItem)
+}
+
 func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
+	// WeChat Test Account
+	// The fields of cleanUser may be incomplete!
 	if cleanUser.WeChatTestAccountId == originUser.WeChatTestAccountId {
 		cleanUser.WeChatTestAccountId = ""
 	}
 	if cleanUser.WeChatTestAccountSecret == originUser.WeChatTestAccountSecret {
 		cleanUser.WeChatTestAccountSecret = ""
 	}
+	// This means the user updated those fields.
 	if cleanUser.WeChatTestAccountId != "" || cleanUser.WeChatTestAccountSecret != "" {
 		oldWeChatTestAccountTokenStoreItem := WeChatTestAccountTokenStoreItem{
 			AppID:     originUser.WeChatTestAccountId,
 			AppSecret: originUser.WeChatTestAccountSecret,
 		}
+		// Yeah, it's a deep copy.
 		newWeChatTestAccountTokenStoreItem := oldWeChatTestAccountTokenStoreItem
 		if cleanUser.WeChatTestAccountId != "" {
 			newWeChatTestAccountTokenStoreItem.AppID = cleanUser.WeChatTestAccountId
@@ -115,6 +138,8 @@ func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
 		}
 		TokenStoreAddItem(&newWeChatTestAccountTokenStoreItem)
 	}
+
+	// WeChat Corp Account
 	if cleanUser.WeChatCorpAccountId == originUser.WeChatCorpAccountId {
 		cleanUser.WeChatCorpAccountId = ""
 	}
@@ -126,9 +151,9 @@ func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
 	}
 	if cleanUser.WeChatCorpAccountId != "" || cleanUser.WeChatCorpAccountAgentId != "" || cleanUser.WeChatCorpAccountAgentSecret != "" {
 		oldWeChatCorpAccountTokenStoreItem := WeChatCorpAccountTokenStoreItem{
-			CorpId:      cleanUser.WeChatCorpAccountId,
+			CorpId:      originUser.WeChatCorpAccountId,
-			AgentSecret: cleanUser.WeChatCorpAccountAgentSecret,
+			AgentSecret: originUser.WeChatCorpAccountAgentSecret,
-			AgentId:     cleanUser.WeChatCorpAccountAgentId,
+			AgentId:     originUser.WeChatCorpAccountAgentId,
 		}
 		newWeChatCorpAccountTokenStoreItem := oldWeChatCorpAccountTokenStoreItem
 		if cleanUser.WeChatCorpAccountId != "" {
@@ -147,7 +172,9 @@ func TokenStoreUpdateUser(cleanUser *model.User, originUser *model.User) {
 	}
 }
 
-// TokenStoreRemoveUser user must be filled
+// TokenStoreRemoveUser
+// user must be filled.
+// It's okay to delete a user that don't have an item here.
 func TokenStoreRemoveUser(user *model.User) {
 	testAccountTokenStoreItem := WeChatTestAccountTokenStoreItem{
 		AppID:     user.WeChatTestAccountId,

channel/wechat-corp-account.go

@@ -34,6 +34,10 @@ func (i *WeChatCorpAccountTokenStoreItem) IsShared() bool {
 		i.CorpId, i.AgentSecret, i.AgentId).Find(&model.User{}).RowsAffected != 1
 }
 
+func (i *WeChatCorpAccountTokenStoreItem) IsFilled() bool {
+	return i.CorpId != "" && i.AgentSecret != "" && i.AgentId != ""
+}
+
 func (i *WeChatCorpAccountTokenStoreItem) Token() string {
 	return i.AccessToken
 }

channel/wechat-test-account.go

@@ -33,6 +33,10 @@ func (i *WeChatTestAccountTokenStoreItem) IsShared() bool {
 		i.AppID, i.AppSecret).Find(&model.User{}).RowsAffected != 1
 }
 
+func (i *WeChatTestAccountTokenStoreItem) IsFilled() bool {
+	return i.AppID != "" && i.AppSecret != ""
+}
+
 func (i *WeChatTestAccountTokenStoreItem) Token() string {
 	return i.AccessToken
 }

controller/user.go

@@ -584,8 +584,24 @@ func ManageUser(c *gin.Context) {
 	}
 	switch req.Action {
 	case "disable":
+		if user.Status == common.UserStatusDisabled {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "该账户已经是封禁状态",
+			})
+			return
+		}
+		channel.TokenStoreRemoveUser(&user)
 		user.Status = common.UserStatusDisabled
 	case "enable":
+		if user.Status == common.UserStatusEnabled {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "该账户已经是启用状态",
+			})
+			return
+		}
+		channel.TokenStoreAddUser(&user)
 		user.Status = common.UserStatusEnabled
 	case "delete":
 		if err := user.Delete(); err != nil {

model/user.go

@@ -53,7 +53,7 @@ func GetAllUsers(startIdx int, num int) (users []*User, err error) {
 }
 
 func GetAllUsersWithSecrets() (users []*User, err error) {
-	err = DB.Find(&users).Error
+	err = DB.Where("status = ?", common.UserStatusEnabled).Where("wechat_test_account_id != '' or wechat_corp_account_id != ''").Find(&users).Error
 	return users, err
 }