light-apiengine-feature/utils/token_store.go

package utils

import (
	"errors"
	"log"
	"strings"
	"sync"
	"time"

	"git.qianqiusoft.com/qianqiusoft/light-apiengine/config"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/entitys"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/logs"
)

type TokenStore interface {
	Get(string) *entitys.Token
	Set(string, *entitys.Token)
	Remove(key string)
	Refresh(key string)
}

type MemoryStore struct {
	name   string
	lock   *sync.RWMutex
	tokens map[string]*entitys.Token
}

var globalTokenStore TokenStore

func init() {
	memoryStore := &MemoryStore{name: "sso", lock: new(sync.RWMutex), tokens: make(map[string]*entitys.Token)}
	globalTokenStore = memoryStore
	go memoryStore.startTokenCheckProcess()
}

func GetGlobalTokenStore() TokenStore {
	return globalTokenStore
}

func SetGlobalTokenStore(store TokenStore) {
	globalTokenStore = store
}

func (t *MemoryStore) Get(key string) *entitys.Token {
	t.lock.RLock()
	defer t.lock.RUnlock()
	if val, ok := t.tokens[key]; ok {
		//log.Println(key, "获取Token:", val.AccessToken, val.RefreshToken, val.LoginID)
		return val
	}
	return nil
}

func (t *MemoryStore) Set(key string, v *entitys.Token) {
	t.lock.Lock()
	defer t.lock.Unlock()
	if val, ok := t.tokens[key]; !ok {
		t.tokens[key] = v
	} else if val != v {
		t.tokens[key] = v
	}
	log.Println(key, "添加Token:", v.AccessToken, v.RefreshToken, v.LoginID)
}

func (t *MemoryStore) Remove(key string) {
	t.lock.Lock()
	defer t.lock.Unlock()
	delete(t.tokens, key)
	log.Println(key, "删除Key")
}

func (t *MemoryStore) Refresh(key string) {
	t.lock.Lock()
	defer t.lock.Unlock()

	val, ok := t.tokens[key]
	if ok {
		val.TimeStamp = uint64(time.Now().UnixNano())
	}
}

func (t *MemoryStore) startTokenCheckProcess() {
	autoRefresh := config.AppConfig.AutoRefresh
	if !autoRefresh {
		return
	}
	var duration time.Duration = time.Second * 3600
	t1 := time.NewTicker(duration)
	for {
		select {
		case <-t1.C:
			t.lock.Lock()

			keys := []string{}
			for k, v := range t.tokens {
				timestampt := time.Unix(int64(v.TimeStamp), 0)
				subval := time.Now().Sub(timestampt)

				if subval.Seconds() < 0 || subval.Seconds() > 3600 {
					keys = append(keys, k)
				}
			}
			for _, k := range keys {
				delete(t.tokens, k)
			}

			t.lock.Unlock()
		}
	}
}

func Validate(accessToken, loginId string, domain string) (*entitys.Token, error) {
	token := globalTokenStore.Get(loginId + domain)
	if token != nil {
		if strings.EqualFold(token.AccessToken, accessToken) {
			logs.Info("get the token ", accessToken, " of id ", loginId+domain)
			globalTokenStore.Refresh(loginId + domain)
			return token, nil
		} else {
			logs.Error(token.AccessToken, "is not equal to", accessToken)
			return token, errors.New(token.AccessToken + " is not equal to " + accessToken)
		}
	} else {
		logs.Error("can not get the token of", loginId+domain)
		return token, errors.New("can not get the token of " + loginId + domain)
	}

}

func TokenValidate(token string) (*entitys.Token, error) {
	user := globalTokenStore.Get(token)

	if strings.EqualFold(user.AccessToken, token) {
		logs.Info("get the token ", token, " of id ")
		globalTokenStore.Refresh(token)
		return user, nil
	} else {
		logs.Error(user.AccessToken, "is not equal to", token)
		return user, errors.New(user.AccessToken + " is not equal to " + token)
	}

}