Tolan/backend/session_store.py

import asyncio
import logging
import time
from dataclasses import dataclass
from typing import Awaitable, Callable, Dict, Optional

from fastapi import WebSocket

logger = logging.getLogger(__name__)


@dataclass
class SessionEntry:
    token: str
    expires_at: float  # unix timestamp (seconds)


class SessionStore:
    """
    简单的会话存储，支持单用户单活 token。
    - 多进程（gunicorn workers）下必须以 Redis 为权威来源；内存仅作缓存。
    - Redis 不可用时回退为单进程内存模式（仅适用于单 worker）。
    """

    def __init__(
        self,
        default_ttl_seconds: int = 1800,
        redis_getter: Optional[Callable[[], Awaitable[object]]] = None,
    ) -> None:
        self._sessions: Dict[str, SessionEntry] = {}
        self._connections: Dict[str, WebSocket] = {}
        self._default_ttl = default_ttl_seconds
        # 保护连接映射的锁，避免并发覆盖
        self._lock = asyncio.Lock()
        # 获取 Redis 客户端的异步函数（来自 DatabaseManager.get_redis）
        self._redis_getter = redis_getter

    def _now(self) -> float:
        return time.time()

    async def _get_redis(self):
        """获取 Redis 客户端；如果未配置或出错则返回 None。"""
        if not self._redis_getter:
            return None
        try:
            redis = await self._redis_getter()
            return redis
        except Exception:
            # Redis 故障时退回内存模式，不影响主流程
            return None

    def _redis_key(self, user_id: str) -> str:
        return f"session:{user_id}"

    async def set_current_token(
        self,
        user_id: str,
        token: str,
        ttl_seconds: Optional[int] = None,
    ) -> None:
        """设置当前用户的会话 token，写入 Redis 并更新本地缓存。"""
        ttl = ttl_seconds or self._default_ttl
        expires_at = self._now() + ttl
        self._sessions[user_id] = SessionEntry(token=token, expires_at=expires_at)

        redis = await self._get_redis()
        if redis:
            try:
                await redis.set(self._redis_key(user_id), token, ex=ttl)
            except Exception as e:
                logger.warning(
                    "session_store: Redis SET failed user_id=%s (multi-worker needs Redis): %s",
                    user_id,
                    e,
                )

    async def get_current_token(self, user_id: str) -> Optional[SessionEntry]:
        """
        获取当前用户的会话信息。
        Redis 可用时以 Redis 为准并回填本地缓存，保证多 worker 间 token 一致。
        """
        now = self._now()
        redis = await self._get_redis()

        if redis:
            try:
                key = self._redis_key(user_id)
                token = await redis.get(key)
                if not token:
                    self._sessions.pop(user_id, None)
                    self._connections.pop(user_id, None)
                    return None

                ttl = await redis.ttl(key)
                # -2: 不存在（竞态）
                if ttl == -2:
                    self._sessions.pop(user_id, None)
                    return None
                # -1: 未设置过期时间，用默认 TTL 估算本地过期时间
                if ttl == -1:
                    expires_at = now + float(self._default_ttl)
                elif ttl is None or ttl <= 0:
                    try:
                        await redis.delete(key)
                    except Exception:
                        pass
                    self._sessions.pop(user_id, None)
                    self._connections.pop(user_id, None)
                    return None
                else:
                    expires_at = now + float(ttl)

                entry = SessionEntry(token=token, expires_at=expires_at)
                self._sessions[user_id] = entry
                return entry
            except Exception as e:
                logger.warning(
                    "session_store: Redis read failed user_id=%s, fallback to memory: %s",
                    user_id,
                    e,
                )

        entry = self._sessions.get(user_id)
        if not entry:
            return None
        if entry.expires_at < now:
            self._sessions.pop(user_id, None)
            self._connections.pop(user_id, None)
            redis = await self._get_redis()
            if redis:
                try:
                    await redis.delete(self._redis_key(user_id))
                except Exception:
                    pass
            return None
        return entry

    async def validate(
        self,
        user_id: Optional[str],
        token: Optional[str],
        refresh: bool = False,
    ) -> bool:
        """校验 token 是否有效，可选滑动续期（刷新 TTL）。"""
        if not user_id or not token:
            return False

        entry = await self.get_current_token(user_id)
        if not entry:
            return False
        if entry.token != token:
            return False

        if refresh:
            # 滑动续期：刷新本地过期时间和 Redis TTL
            now = self._now()
            entry.expires_at = now + self._default_ttl
            self._sessions[user_id] = entry

            redis = await self._get_redis()
            if redis:
                try:
                    await redis.expire(self._redis_key(user_id), self._default_ttl)
                except Exception:
                    pass

        return True

    async def bind_connection(self, user_id: str, token: str, websocket: WebSocket) -> None:
        """
        记录当前生效 token 对应的连接，用于后续定向下线。
        只有当 token 等于当前存储的 token 时才绑定。
        """
        async with self._lock:
            entry = await self.get_current_token(user_id)
            if not entry or entry.token != token:
                return
            # 关闭旧连接
            prev_conn = self._connections.get(user_id)
            if prev_conn and prev_conn is not websocket:
                try:
                    await prev_conn.close(code=4401, reason="session_invalidated")
                except Exception:
                    pass
            self._connections[user_id] = websocket

    async def invalidate_connection(self, user_id: str) -> None:
        """
        主动关闭某个用户的现有连接（用于新 token 签发后踢旧连）。
        """
        async with self._lock:
            conn = self._connections.get(user_id)
            if conn:
                try:
                    await conn.close(code=4401, reason="session_invalidated")
                except Exception:
                    pass
                self._connections.pop(user_id, None)