117 lines
2.3 KiB
Go
117 lines
2.3 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TokenStore 刷新令牌存储接口
|
|
type TokenStore interface {
|
|
// Setex 存储 refresh token，过期时间 seconds
|
|
Setex(ctx context.Context, key string, value string, seconds int) error
|
|
// Get 获取 refresh token
|
|
Get(ctx context.Context, key string) (string, error)
|
|
// Del 删除 refresh token（退出登录）
|
|
Del(ctx context.Context, key string) error
|
|
}
|
|
|
|
// ErrNotFound 表示 key 不存在
|
|
var ErrNotFound = errors.New("token not found")
|
|
|
|
// InMemoryTokenStore 内存存储刷新令牌，用于单机测试
|
|
// 不需要 Redis 依赖
|
|
type InMemoryTokenStore struct {
|
|
mu sync.RWMutex
|
|
data map[string]tokenEntry
|
|
stop chan struct{}
|
|
timer *time.Ticker
|
|
}
|
|
|
|
type tokenEntry struct {
|
|
value string
|
|
expires time.Time
|
|
}
|
|
|
|
// NewInMemoryTokenStore 创建内存令牌存储
|
|
// 每 60 秒自动清理过期项
|
|
func NewInMemoryTokenStore() *InMemoryTokenStore {
|
|
store := &InMemoryTokenStore{
|
|
data: make(map[string]tokenEntry),
|
|
stop: make(chan struct{}),
|
|
}
|
|
store.timer = time.NewTicker(60 * time.Second)
|
|
go store.cleanupLoop()
|
|
return store
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) Setex(_ context.Context, key string, value string, seconds int) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
s.data[key] = tokenEntry{
|
|
value: value,
|
|
expires: time.Now().Add(time.Duration(seconds) * time.Second),
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) Del(_ context.Context, key string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
delete(s.data, key)
|
|
return nil
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) Get(_ context.Context, key string) (string, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
entry, ok := s.data[key]
|
|
if !ok {
|
|
return "", ErrNotFound
|
|
}
|
|
if time.Now().After(entry.expires) {
|
|
delete(s.data, key)
|
|
return "", ErrNotFound
|
|
}
|
|
return entry.value, nil
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) cleanupLoop() {
|
|
for {
|
|
select {
|
|
case <-s.timer.C:
|
|
func() {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
slog.Error("token存储: 清理panic恢复", "panic", r)
|
|
}
|
|
}()
|
|
s.cleanup()
|
|
}()
|
|
case <-s.stop:
|
|
s.timer.Stop()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) cleanup() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
for k, v := range s.data {
|
|
if now.After(v.expires) {
|
|
delete(s.data, k)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *InMemoryTokenStore) Stop() {
|
|
close(s.stop)
|
|
}
|