Session 隔離模式

# sessions/model.py
"""
具密碼學邊界的 Session 模型。
"""
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional
import secrets
import hashlib
 
 
@dataclass
class IsolatedSession:
    session_id: str
    user_id: str
    created_at: datetime
    token: str  # 密碼學 session token
    context: list[dict] = field(default_factory=list)
    metadata: dict = field(default_factory=dict)
    max_context_length: int = 50
    is_active: bool = True
 
    @staticmethod
    def create(user_id: str, max_context: int = 50) -> "IsolatedSession":
        session_id = secrets.token_hex(16)
        token = secrets.token_urlsafe(32)
        return IsolatedSession(
            session_id=session_id,
            user_id=user_id,
            created_at=datetime.now(timezone.utc),
            token=token,
            max_context_length=max_context,
        )
 
    def add_message(self, role: str, content: str) -> None:
        self.context.append({"role": role, "content": content})
        if len(self.context) > self.max_context_length:
            self.context = self.context[-self.max_context_length:]
 
    def get_context(self) -> list[dict]:
        return list(self.context)

# sessions/store.py
"""
強制執行邊界的隔離 Session 儲存。
"""
import logging
from typing import Optional
from sessions.model import IsolatedSession
 
logger = logging.getLogger("session_store")
 
 
class SessionStore:
    def __init__(self):
        self._sessions: dict[str, IsolatedSession] = {}
        self._user_sessions: dict[str, set[str]] = {}
 
    def create_session(self, user_id: str) -> IsolatedSession:
        session = IsolatedSession.create(user_id)
        self._sessions[session.session_id] = session
        if user_id not in self._user_sessions:
            self._user_sessions[user_id] = set()
        self._user_sessions[user_id].add(session.session_id)
        return session
 
    def get_session(
        self, session_id: str, token: str
    ) -> Optional[IsolatedSession]:
        session = self._sessions.get(session_id)
        if not session:
            return None
        if not secrets.compare_digest(session.token, token):
            logger.warning(f"Invalid token for session {session_id}")
            return None
        if not session.is_active:
            return None
        return session
 
    def terminate_session(self, session_id: str) -> None:
        session = self._sessions.get(session_id)
        if session:
            session.is_active = False
            session.context.clear()
            logger.info(f"Terminated session {session_id}")
 
    def get_user_sessions(self, user_id: str) -> list[str]:
        return list(self._user_sessions.get(user_id, set()))

# sessions/isolation.py
"""
上下文隔離強制執行。
"""
import hashlib
 
 
class ContextIsolationEnforcer:
    def __init__(self, store):
        self.store = store
 
    def validate_context_request(
        self, requesting_session_id: str, target_session_id: str
    ) -> bool:
        if requesting_session_id == target_session_id:
            return True
        return False  # 絕不允許跨 session 上下文存取
 
    def sanitize_shared_resources(
        self, resource: str, session_id: str
    ) -> str:
        """確保共用資源不含 session 特定資料。"""
        session = self.store._sessions.get(session_id)
        if not session:
            return resource
        for msg in session.context:
            if msg["content"] in resource:
                resource = resource.replace(
                    msg["content"], "[REDACTED_SESSION_DATA]"
                )
        return resource

# sessions/memory.py
"""
LLM session 的隔離記憶體管理。
"""
from typing import Optional
 
 
class SessionMemoryManager:
    def __init__(self):
        self._memories: dict[str, dict[str, str]] = {}
 
    def store(self, session_id: str, key: str, value: str) -> None:
        if session_id not in self._memories:
            self._memories[session_id] = {}
        self._memories[session_id][key] = value
 
    def retrieve(self, session_id: str, key: str) -> Optional[str]:
        return self._memories.get(session_id, {}).get(key)
 
    def clear_session(self, session_id: str) -> None:
        self._memories.pop(session_id, None)
 
    def list_keys(self, session_id: str) -> list[str]:
        return list(self._memories.get(session_id, {}).keys())

# sessions/api.py
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from sessions.store import SessionStore
 
app = FastAPI(title="Session Isolation Service")
store = SessionStore()
 
 
class CreateSessionRequest(BaseModel):
    user_id: str
 
 
class MessageRequest(BaseModel):
    message: str
 
 
@app.post("/session/create")
async def create_session(request: CreateSessionRequest):
    session = store.create_session(request.user_id)
    return {"session_id": session.session_id, "token": session.token}
 
 
@app.post("/session/{session_id}/message")
async def send_message(
    session_id: str, request: MessageRequest,
    x_session_token: str = Header(...),
):
    session = store.get_session(session_id, x_session_token)
    if not session:
        raise HTTPException(403, "Invalid session")
    session.add_message("user", request.message)
    return {"context_length": len(session.context)}

uvicorn sessions.api:app --port 8650

# tests/test_sessions.py
import pytest
from sessions.store import SessionStore
 
 
@pytest.fixture
def store():
    return SessionStore()
 
def test_sessions_are_isolated(store):
    s1 = store.create_session("user-1")
    s2 = store.create_session("user-2")
    s1.add_message("user", "Secret information for user 1")
    assert "Secret" not in str(s2.get_context())
 
def test_invalid_token_rejected(store):
    s1 = store.create_session("user-1")
    result = store.get_session(s1.session_id, "wrong-token")
    assert result is None
 
def test_terminated_session_inaccessible(store):
    s1 = store.create_session("user-1")
    store.terminate_session(s1.session_id)
    result = store.get_session(s1.session_id, s1.token)
    assert result is None

pytest tests/test_sessions.py -v

# sessions/monitoring.py
import logging
from collections import Counter
 
 
class SessionMonitor:
    def __init__(self):
        self.invalid_token_attempts = Counter()
        self.logger = logging.getLogger("session_monitor")
 
    def record_invalid_token(self, session_id: str, source_ip: str):
        self.invalid_token_attempts[source_ip] += 1
        if self.invalid_token_attempts[source_ip] >= 5:
            self.logger.critical(
                f"Possible session hijack attempt from {source_ip}: "
                f"{self.invalid_token_attempts[source_ip]} invalid tokens"
            )