light-apiengine-feature/utils/store.go

package utils

import (
	"log"
	"sync"
	"time"

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

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

var GlobalTokenStore *TokenStore

func NewTokenStore() *TokenStore {
	t := &TokenStore{name: "sso", lock: new(sync.RWMutex), tokens: make(map[string]*entitys.Token)}
	go t.startTokenCheckProcess()
	return t
}

func (t *TokenStore) 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 *TokenStore) 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 *TokenStore) Remove(key string) {
	t.lock.Lock()
	defer t.lock.Unlock()
	delete(t.tokens, key)
	log.Println(key, "删除Key")
}

func (t *TokenStore) 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 *TokenStore) startTokenCheckProcess() {
	autoRefresh, _ := NewConfig().Bool("auto_refresh", false)
	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()
		}
	}
}