geo/server/internal/shared/auth/session_store.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

var (
	ErrRefreshSessionNotFound = errors.New("refresh session not found")
	ErrRefreshTokenMismatch   = errors.New("refresh token mismatch")
)

var rotateRefreshScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
if not current then
  return 0
end
if current == ARGV[1] then
  redis.call("PSETEX", KEYS[1], ARGV[4], "grace:" .. ARGV[1])
  redis.call("PSETEX", KEYS[2], ARGV[3], ARGV[2])
  return 1
end
if current == "grace:" .. ARGV[1] then
  redis.call("PSETEX", KEYS[2], ARGV[3], ARGV[2])
  return 1
end
return -1
`)

var setMaxSessionVersionScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
local current_number = tonumber(current or "0") or 0
local next_number = tonumber(ARGV[1]) or 0
if current_number >= next_number then
  return current_number
end
redis.call("SET", KEYS[1], ARGV[1])
return next_number
`)

const sessionVersionFallbackTTL = 30 * 24 * time.Hour

// refreshRotationGraceTTL keeps a rotated refresh token usable for a short
// window so concurrent refreshes (e.g. multiple browser tabs racing with the
// same token) each obtain a valid new pair instead of being logged out.
const refreshRotationGraceTTL = 60 * time.Second

const rotationGracePrefix = "grace:"

type SessionStore struct {
	rdb      *redis.Client
	fallback *memorySessionStore
}

func NewSessionStore(rdb *redis.Client) *SessionStore {
	return &SessionStore{
		rdb:      rdb,
		fallback: newMemorySessionStore(),
	}
}

func (s *SessionStore) SaveRefresh(ctx context.Context, jti string, tokenHash string, ttl time.Duration) error {
	key := fmt.Sprintf("refresh:%s", jti)
	s.fallback.set(key, tokenHash, ttl)
	if s.rdb == nil {
		return nil
	}
	if err := s.rdb.Set(ctx, key, tokenHash, ttl).Err(); err != nil {
		return nil
	}
	return nil
}

func (s *SessionStore) GetRefresh(ctx context.Context, jti string) (string, error) {
	key := fmt.Sprintf("refresh:%s", jti)
	if s.rdb != nil {
		value, err := s.rdb.Get(ctx, key).Result()
		if err == nil {
			if strings.HasPrefix(value, rotationGracePrefix) {
				return "", ErrRefreshSessionNotFound
			}
			return value, nil
		}
	}
	value, err := s.fallback.get(key)
	if err != nil {
		return "", err
	}
	if strings.HasPrefix(value, rotationGracePrefix) {
		return "", ErrRefreshSessionNotFound
	}
	return value, nil
}

func (s *SessionStore) DeleteRefresh(ctx context.Context, jti string) error {
	key := fmt.Sprintf("refresh:%s", jti)
	s.fallback.delete(key)
	if s.rdb == nil {
		return nil
	}
	if err := s.rdb.Del(ctx, key).Err(); err != nil {
		return nil
	}
	return nil
}

func (s *SessionStore) RotateRefresh(ctx context.Context, oldJTI, oldHash, newJTI, newHash string, ttl time.Duration) error {
	oldKey := fmt.Sprintf("refresh:%s", oldJTI)
	newKey := fmt.Sprintf("refresh:%s", newJTI)

	if s.rdb == nil {
		return s.fallback.rotate(oldKey, oldHash, newKey, newHash, ttl)
	}

	result, err := rotateRefreshScript.Run(
		ctx,
		s.rdb,
		[]string{oldKey, newKey},
		oldHash,
		newHash,
		ttl.Milliseconds(),
		refreshRotationGraceTTL.Milliseconds(),
	).Int()
	if err != nil {
		return s.fallback.rotate(oldKey, oldHash, newKey, newHash, ttl)
	}

	switch result {
	case 1:
		_ = s.fallback.rotate(oldKey, oldHash, newKey, newHash, ttl)
		return nil
	case 0:
		if err := s.fallback.rotate(oldKey, oldHash, newKey, newHash, ttl); err == nil {
			return nil
		}
		return ErrRefreshSessionNotFound
	case -1:
		return ErrRefreshTokenMismatch
	default:
		return fmt.Errorf("unexpected rotate refresh result: %d", result)
	}
}

func (s *SessionStore) Blacklist(ctx context.Context, accessJTI string, ttl time.Duration) error {
	key := fmt.Sprintf("blacklist:%s", accessJTI)
	s.fallback.set(key, "1", ttl)
	if s.rdb == nil {
		return nil
	}
	if err := s.rdb.Set(ctx, key, "1", ttl).Err(); err != nil {
		return nil
	}
	return nil
}

func (s *SessionStore) IsBlacklisted(ctx context.Context, accessJTI string) (bool, error) {
	key := fmt.Sprintf("blacklist:%s", accessJTI)
	if _, err := s.fallback.get(key); err == nil {
		return true, nil
	}
	if s.rdb == nil {
		return false, nil
	}
	n, err := s.rdb.Exists(ctx, key).Result()
	if err != nil {
		return false, nil
	}
	return n > 0, nil
}

func (s *SessionStore) SessionVersion(ctx context.Context, userID int64) (int64, error) {
	if s == nil {
		return 0, nil
	}
	key := sessionVersionKey(userID)
	fallbackVersion := s.fallback.sessionVersion(key)
	redisVersion := int64(0)
	if s.rdb != nil {
		value, err := s.rdb.Get(ctx, key).Result()
		if err == nil {
			version, parseErr := strconv.ParseInt(value, 10, 64)
			if parseErr == nil && version > 0 {
				redisVersion = version
			}
		}
	}

	version := maxInt64(redisVersion, fallbackVersion)
	if version <= 0 {
		return 0, nil
	}
	s.fallback.set(key, strconv.FormatInt(version, 10), sessionVersionFallbackTTL)
	if s.rdb != nil && version > redisVersion {
		_, _ = setMaxSessionVersionScript.Run(ctx, s.rdb, []string{key}, strconv.FormatInt(version, 10)).Int64()
	}
	return version, nil
}

func (s *SessionStore) BumpSessionVersion(ctx context.Context, userID int64) (int64, error) {
	if s == nil {
		return 0, nil
	}
	key := sessionVersionKey(userID)
	current, _ := s.SessionVersion(ctx, userID)
	next := current + 1
	s.fallback.set(key, strconv.FormatInt(next, 10), sessionVersionFallbackTTL)

	if s.rdb == nil {
		return next, nil
	}
	version, err := s.rdb.Incr(ctx, key).Result()
	if err != nil {
		return next, nil
	}
	if version < next {
		if maxVersion, maxErr := setMaxSessionVersionScript.Run(ctx, s.rdb, []string{key}, strconv.FormatInt(next, 10)).Int64(); maxErr == nil {
			version = maxVersion
		}
	}
	s.fallback.set(key, strconv.FormatInt(version, 10), sessionVersionFallbackTTL)
	return version, nil
}

func sessionVersionKey(userID int64) string {
	return fmt.Sprintf("session_version:user:%d", userID)
}

func maxInt64(a, b int64) int64 {
	if a > b {
		return a
	}
	return b
}

type memorySessionStore struct {
	mu    sync.Mutex
	items map[string]memorySessionItem
}

type memorySessionItem struct {
	value     string
	expiresAt time.Time
}

func newMemorySessionStore() *memorySessionStore {
	return &memorySessionStore{items: make(map[string]memorySessionItem)}
}

func (s *memorySessionStore) set(key, value string, ttl time.Duration) {
	if s == nil || ttl <= 0 {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.items[key] = memorySessionItem{
		value:     value,
		expiresAt: time.Now().Add(ttl),
	}
}

func (s *memorySessionStore) get(key string) (string, error) {
	if s == nil {
		return "", redis.Nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	item, ok := s.items[key]
	if !ok {
		return "", redis.Nil
	}
	if !item.expiresAt.After(time.Now()) {
		delete(s.items, key)
		return "", redis.Nil
	}
	return item.value, nil
}

func (s *memorySessionStore) delete(key string) {
	if s == nil {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.items, key)
}

func (s *memorySessionStore) sessionVersion(key string) int64 {
	if s == nil {
		return 0
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	if item, ok := s.items[key]; ok && item.expiresAt.After(time.Now()) {
		if parsed, err := strconv.ParseInt(item.value, 10, 64); err == nil && parsed > 0 {
			return parsed
		}
	}
	return 0
}

func (s *memorySessionStore) rotate(oldKey, oldHash, newKey, newHash string, ttl time.Duration) error {
	if s == nil {
		return ErrRefreshSessionNotFound
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	item, ok := s.items[oldKey]
	if !ok {
		return ErrRefreshSessionNotFound
	}
	if !item.expiresAt.After(time.Now()) {
		delete(s.items, oldKey)
		return ErrRefreshSessionNotFound
	}
	switch item.value {
	case oldHash:
		s.items[oldKey] = memorySessionItem{
			value:     rotationGracePrefix + oldHash,
			expiresAt: time.Now().Add(refreshRotationGraceTTL),
		}
	case rotationGracePrefix + oldHash:
		// Already rotated within the grace window; keep the marker as-is.
	default:
		return ErrRefreshTokenMismatch
	}
	if ttl > 0 {
		s.items[newKey] = memorySessionItem{
			value:     newHash,
			expiresAt: time.Now().Add(ttl),
		}
	}
	return nil
}