Files
Kiro/internal/auth/state_store.go

79 lines
1.6 KiB
Go

package auth
import (
"sync"
"time"
)
// MemoryStateStore is an in-memory implementation of StateStore
// Note: This is suitable for development but should be replaced with
// a distributed store (Redis, DynamoDB) for production with multiple servers
type MemoryStateStore struct {
mu sync.RWMutex
states map[string]time.Time
}
// NewMemoryStateStore creates a new in-memory state store
func NewMemoryStateStore() *MemoryStateStore {
store := &MemoryStateStore{
states: make(map[string]time.Time),
}
// Start cleanup goroutine to remove expired states
go store.cleanupExpired()
return store
}
// Set stores a state token with an expiry time
func (s *MemoryStateStore) Set(state string, expiry time.Time) error {
s.mu.Lock()
defer s.mu.Unlock()
s.states[state] = expiry
return nil
}
// Validate checks if a state token is valid and not expired
func (s *MemoryStateStore) Validate(state string) (bool, error) {
s.mu.RLock()
defer s.mu.RUnlock()
expiry, exists := s.states[state]
if !exists {
return false, nil
}
if time.Now().After(expiry) {
return false, nil
}
return true, nil
}
// Delete removes a state token from the store
func (s *MemoryStateStore) Delete(state string) error {
s.mu.Lock()
defer s.mu.Unlock()
delete(s.states, state)
return nil
}
// cleanupExpired periodically removes expired state tokens
func (s *MemoryStateStore) cleanupExpired() {
ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop()
for range ticker.C {
s.mu.Lock()
now := time.Now()
for state, expiry := range s.states {
if now.After(expiry) {
delete(s.states, state)
}
}
s.mu.Unlock()
}
}