Tolan/backend/session_store.py

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

from fastapi import WebSocket


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


class SessionStore:
    """
    简单的会话存储，支持单用户单活 token。
    - 默认使用内存字典缓存当前进程内活跃会话
    - 同步使用 Redis 持久化 token，支持进程重启后继续校验
    """

    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:
            # 使用 Redis 仅存储 token，过期时间通过 TTL 控制
            try:
                await redis.set(self._redis_key(user_id), token, ex=ttl)
            except Exception:
                # Redis 写入失败时不影响当前进程内会话
                pass

    async def get_current_token(self, user_id: str) -> Optional[SessionEntry]:
        """
        获取当前用户的会话信息。
        优先从内存缓存读取；若不存在则尝试从 Redis 恢复。
        """
        entry = self._sessions.get(user_id)
        now = self._now()

        if entry:
            if entry.expires_at < now:
                # 已过期，清理本地和 Redis
                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

        # 内存中没有，尝试从 Redis 恢复（支持进程重启）
        redis = await self._get_redis()
        if not redis:
            return None

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

            ttl = await redis.ttl(key)
            # ttl: 剩余秒数，-1=无过期，-2=不存在
            if ttl is None or ttl <= 0:
                await redis.delete(key)
                return None

            entry = SessionEntry(token=token, expires_at=now + ttl)
            self._sessions[user_id] = entry
            return entry
        except Exception:
            # Redis 故障时直接视为无会话
            return None

    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)