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
|
|||
|
|
}
|