byop-engine/auth/memory_ts.go

package auth

import (
	"context"
	"sync"
	"time"
)

// MemoryTokenStore implements TokenStore with in-memory storage
type MemoryTokenStore struct {
	blacklist map[string]time.Time
	mutex     sync.RWMutex
	// Optional cleaner to remove expired tokens
	cleanupInterval time.Duration
	stopCleanup     chan struct{}
}

// NewMemoryTokenStore creates a new in-memory token store
func NewMemoryTokenStore(cleanupInterval time.Duration) *MemoryTokenStore {
	store := &MemoryTokenStore{
		blacklist:       make(map[string]time.Time),
		mutex:           sync.RWMutex{},
		cleanupInterval: cleanupInterval,
		stopCleanup:     make(chan struct{}),
	}

	// Start background cleanup if interval > 0
	if cleanupInterval > 0 {
		go store.startCleanupRoutine()
	}

	return store
}

// IsBlacklisted checks if a token is in the blacklist
func (s *MemoryTokenStore) IsBlacklisted(ctx context.Context, token string) (bool, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	expiry, exists := s.blacklist[token]
	if !exists {
		return false, nil
	}

	// If token has expired, we can remove it
	if time.Now().After(expiry) {
		// Don't remove here to avoid write lock during read operation
		return false, nil
	}

	return true, nil
}

// Blacklist adds a token to the blacklist
func (s *MemoryTokenStore) Blacklist(ctx context.Context, token string, expiry time.Time) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.blacklist[token] = expiry
	return nil
}

// startCleanupRoutine periodically removes expired tokens
func (s *MemoryTokenStore) startCleanupRoutine() {
	ticker := time.NewTicker(s.cleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.cleanup()
		case <-s.stopCleanup:
			return
		}
	}
}

// cleanup removes expired tokens from the blacklist
func (s *MemoryTokenStore) cleanup() {
	now := time.Now()
	s.mutex.Lock()
	defer s.mutex.Unlock()

	for token, expiry := range s.blacklist {
		if now.After(expiry) {
			delete(s.blacklist, token)
		}
	}
}

// Close stops the cleanup goroutine
func (s *MemoryTokenStore) Close() {
	if s.cleanupInterval > 0 {
		s.stopCleanup <- struct{}{}
	}
}