message-pusher/channel/token-store.go

package channel

import (
	"message-pusher/common"
	"message-pusher/model"
	"sync"
	"time"
)

type TokenStoreItem interface {
	Key() string
	Token() string
	Refresh()
	IsFilled() bool
	IsShared() bool
}

type tokenStore struct {
	Map               map[string]*TokenStoreItem
	Mutex             sync.RWMutex
	ExpirationSeconds int
}

var s tokenStore

func channel2item(channel_ *model.Channel) TokenStoreItem {
	switch channel_.Type {
	case model.TypeWeChatTestAccount:
		item := &WeChatTestAccountTokenStoreItem{
			AppID:     channel_.AppId,
			AppSecret: channel_.Secret,
		}
		return item
	case model.TypeWeChatCorpAccount:
		corpId, agentId, err := parseWechatCorpAccountAppId(channel_.AppId)
		if err != nil {
			common.SysError(err.Error())
			return nil
		}
		item := &WeChatCorpAccountTokenStoreItem{
			CorpId:      corpId,
			AgentSecret: channel_.Secret,
			AgentId:     agentId,
		}
		return item
	case model.TypeLarkApp:
		item := &LarkAppTokenStoreItem{
			AppID:     channel_.AppId,
			AppSecret: channel_.Secret,
		}
		return item
	}
	return nil
}

func channels2items(channels []*model.Channel) []TokenStoreItem {
	var items []TokenStoreItem
	for _, channel_ := range channels {
		item := channel2item(channel_)
		if item != nil {
			items = append(items, item)
		}
	}
	return items
}

func TokenStoreInit() {
	s.Map = make(map[string]*TokenStoreItem)
	// https://developers.weixin.qq.com/doc/offiaccount/Basic_Information/Get_access_token.html
	// https://developer.work.weixin.qq.com/document/path/91039
	s.ExpirationSeconds = 2 * 55 * 60 // 2 hours - 5 minutes
	go func() {
		channels, err := model.GetTokenStoreChannels()
		if err != nil {
			common.FatalLog(err.Error())
		}
		items := channels2items(channels)
		s.Mutex.RLock()
		for i := range items {
			// s.Map[item.Key()] = &item  // This is wrong, you are getting the address of a local variable!
			s.Map[items[i].Key()] = &items[i]
		}
		s.Mutex.RUnlock()
		for {
			s.Mutex.RLock()
			var tmpMap = make(map[string]*TokenStoreItem)
			for k, v := range s.Map {
				tmpMap[k] = v
			}
			s.Mutex.RUnlock()
			for k := range tmpMap {
				(*tmpMap[k]).Refresh()
			}
			s.Mutex.RLock()
			// we shouldn't directly replace the old map with the new map, cause the old map's keys may already change
			for k := range s.Map {
				v, okay := tmpMap[k]
				if okay {
					s.Map[k] = v
				}
			}
			sleepDuration := common.Max(s.ExpirationSeconds, 60)
			s.Mutex.RUnlock()
			time.Sleep(time.Duration(sleepDuration) * time.Second)
		}
	}()
}

// 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
	s.Mutex.RUnlock()
}

func TokenStoreRemoveItem(item TokenStoreItem) {
	s.Mutex.RLock()
	delete(s.Map, item.Key())
	s.Mutex.RUnlock()
}

func TokenStoreAddUser(user *model.User) {
	channels, err := model.GetTokenStoreChannelsByUserId(user.Id)
	if err != nil {
		common.SysError(err.Error())
		return
	}
	items := channels2items(channels)
	for i := range items {
		TokenStoreAddItem(items[i])
	}
}

// TokenStoreRemoveUser
// user must be filled.
// It's okay to delete a user that don't have an item here.
func TokenStoreRemoveUser(user *model.User) {
	channels, err := model.GetTokenStoreChannelsByUserId(user.Id)
	if err != nil {
		common.SysError(err.Error())
		return
	}
	items := channels2items(channels)
	for i := range items {
		if items[i].IsShared() {
			continue
		}
		TokenStoreRemoveItem(items[i])
	}
}

func checkTokenStoreChannelType(channelType string) bool {
	return channelType == model.TypeWeChatTestAccount || channelType == model.TypeWeChatCorpAccount || channelType == model.TypeLarkApp
}

func TokenStoreAddChannel(channel *model.Channel) {
	if !checkTokenStoreChannelType(channel.Type) {
		return
	}
	item := channel2item(channel)
	if item != nil {
		// Do not check IsShared here, cause its useless
		TokenStoreAddItem(item)
	}
}

func TokenStoreRemoveChannel(channel *model.Channel) {
	if !checkTokenStoreChannelType(channel.Type) {
		return
	}
	item := channel2item(channel)
	if item != nil && !item.IsShared() {
		TokenStoreRemoveItem(item)
	}
}

func TokenStoreUpdateChannel(newChannel *model.Channel, oldChannel *model.Channel) {
	if oldChannel.Type != model.TypeWeChatTestAccount && oldChannel.Type != model.TypeWeChatCorpAccount {
		return
	}
	// Why so complicated? Because the given channel maybe incomplete.
	if oldChannel.Type == model.TypeWeChatTestAccount {
		// Only keep changed parts
		if newChannel.AppId == oldChannel.AppId {
			newChannel.AppId = ""
		}
		if newChannel.Secret == oldChannel.Secret {
			newChannel.Secret = ""
		}
		oldItem := WeChatTestAccountTokenStoreItem{
			AppID:     oldChannel.AppId,
			AppSecret: oldChannel.Secret,
		}
		// Yeah, it's a deep copy.
		newItem := oldItem
		// This means the user updated those fields.
		if newChannel.AppId != "" {
			newItem.AppID = newChannel.AppId
		}
		if newChannel.Secret != "" {
			newItem.AppSecret = newChannel.Secret
		}
		if !oldItem.IsShared() {
			TokenStoreRemoveItem(&oldItem)
		}
		TokenStoreAddItem(&newItem)
		return
	}
	if oldChannel.Type == model.TypeWeChatCorpAccount {
		// Only keep changed parts
		if newChannel.AppId == oldChannel.AppId {
			newChannel.AppId = ""
		}
		if newChannel.Secret == oldChannel.Secret {
			newChannel.Secret = ""
		}
		corpId, agentId, err := parseWechatCorpAccountAppId(oldChannel.AppId)
		if err != nil {
			common.SysError(err.Error())
			return
		}
		oldItem := WeChatCorpAccountTokenStoreItem{
			CorpId:      corpId,
			AgentSecret: oldChannel.Secret,
			AgentId:     agentId,
		}
		// Yeah, it's a deep copy.
		newItem := oldItem
		// This means the user updated those fields.
		if newChannel.AppId != "" {
			corpId, agentId, err := parseWechatCorpAccountAppId(oldChannel.AppId)
			if err != nil {
				common.SysError(err.Error())
				return
			}
			newItem.CorpId = corpId
			newItem.AgentId = agentId
		}
		if newChannel.Secret != "" {
			newItem.AgentSecret = newChannel.Secret
		}
		if !oldItem.IsShared() {
			TokenStoreRemoveItem(&oldItem)
		}
		TokenStoreAddItem(&newItem)
		return
	}
}

func TokenStoreGetToken(key string) string {
	s.Mutex.RLock()
	defer s.Mutex.RUnlock()
	item, ok := s.Map[key]
	if ok {
		return (*item).Token()
	}
	common.SysError("token for " + key + " is blank!")
	return ""
}