114 lines
2.4 KiB
Go
114 lines
2.4 KiB
Go
package auth
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// StateManager OAuth状态管理器
|
||
type StateManager struct {
|
||
states map[string]time.Time
|
||
mu sync.RWMutex
|
||
ttl time.Duration
|
||
}
|
||
|
||
var (
|
||
// 全局状态管理器
|
||
stateManager = &StateManager{
|
||
states: make(map[string]time.Time),
|
||
ttl: 10 * time.Minute, // 10分钟过期
|
||
}
|
||
)
|
||
|
||
// Note: GenerateState and ValidateState are defined in oauth_utils.go
|
||
// to avoid duplication, please use those implementations
|
||
|
||
// Store 存储state
|
||
func (sm *StateManager) Store(state string) {
|
||
sm.mu.Lock()
|
||
defer sm.mu.Unlock()
|
||
sm.states[state] = time.Now()
|
||
}
|
||
|
||
// Validate 验证state
|
||
func (sm *StateManager) Validate(state string) bool {
|
||
sm.mu.RLock()
|
||
defer sm.mu.RUnlock()
|
||
|
||
expiredAt, exists := sm.states[state]
|
||
if !exists {
|
||
return false
|
||
}
|
||
|
||
// 检查是否过期
|
||
return time.Now().Before(expiredAt.Add(sm.ttl))
|
||
}
|
||
|
||
// Delete 删除state（使用后删除）
|
||
func (sm *StateManager) Delete(state string) {
|
||
sm.mu.Lock()
|
||
defer sm.mu.Unlock()
|
||
delete(sm.states, state)
|
||
}
|
||
|
||
// Cleanup 清理过期的state
|
||
func (sm *StateManager) Cleanup() {
|
||
sm.mu.Lock()
|
||
defer sm.mu.Unlock()
|
||
|
||
now := time.Now()
|
||
for state, expiredAt := range sm.states {
|
||
if now.After(expiredAt.Add(sm.ttl)) {
|
||
delete(sm.states, state)
|
||
}
|
||
}
|
||
}
|
||
|
||
// StartCleanupRoutine 启动定期清理goroutine
|
||
// stop channel 关闭时，清理goroutine将优雅退出
|
||
func (sm *StateManager) StartCleanupRoutine(stop <-chan struct{}) {
|
||
ticker := time.NewTicker(5 * time.Minute)
|
||
go func() {
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
sm.Cleanup()
|
||
case <-stop:
|
||
ticker.Stop()
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// CleanupRoutineManager 管理清理goroutine的生命周期
|
||
type CleanupRoutineManager struct {
|
||
stopChan chan struct{}
|
||
}
|
||
|
||
var cleanupRoutineManager *CleanupRoutineManager
|
||
|
||
// StartCleanupRoutineWithManager 使用管理器启动清理goroutine
|
||
func StartCleanupRoutineWithManager() {
|
||
if cleanupRoutineManager != nil {
|
||
return // 已经启动
|
||
}
|
||
cleanupRoutineManager = &CleanupRoutineManager{
|
||
stopChan: make(chan struct{}),
|
||
}
|
||
stateManager.StartCleanupRoutine(cleanupRoutineManager.stopChan)
|
||
}
|
||
|
||
// StopCleanupRoutine 停止清理goroutine（用于优雅关闭）
|
||
func StopCleanupRoutine() {
|
||
if cleanupRoutineManager != nil {
|
||
close(cleanupRoutineManager.stopChan)
|
||
cleanupRoutineManager = nil
|
||
}
|
||
}
|
||
|
||
// GetStateManager 获取全局状态管理器
|
||
func GetStateManager() *StateManager {
|
||
return stateManager
|
||
}
|