优化鉴权系统

2026-07-03 21:54:19 +08:00
parent 7e65fadd83
commit 98a92b83fa
11 changed files with 539 additions and 149 deletions
@@ -0,0 +1,160 @@
+from dataclasses import dataclass
+from datetime import UTC, datetime
+import hmac
+import secrets
+from pathlib import Path
+from uuid import uuid4
+
+import aiosqlite
+import bcrypt
+
+
+class UsernameExistsError(Exception):
+    pass
+
+
+class InvalidCredentialsError(Exception):
+    pass
+
+
+class UserStoreError(Exception):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class RegisteredUser:
+    user_id: str
+    api_key: str
+    created_at: datetime
+
+
+@dataclass(frozen=True, slots=True)
+class AuthenticatedUser:
+    user_id: str
+
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS users (
+    user_id        TEXT PRIMARY KEY,
+    username       TEXT UNIQUE NOT NULL,
+    password_hash  TEXT NOT NULL,
+    api_key        TEXT UNIQUE NOT NULL,
+    created_at     TEXT NOT NULL
+);
+"""
+
+
+class UserStore:
+    def __init__(self, db_path: Path) -> None:
+        self.db_path = db_path
+
+    async def init(self) -> None:
+        try:
+            self.db_path.parent.mkdir(parents=True, exist_ok=True)
+            async with aiosqlite.connect(self.db_path) as db:
+                await db.execute("PRAGMA journal_mode=WAL;")
+                await db.execute("PRAGMA busy_timeout=5000;")
+                await db.execute("PRAGMA foreign_keys=ON;")
+                await db.executescript(_SCHEMA)
+                await db.commit()
+        except aiosqlite.Error as exc:
+            raise UserStoreError("could not initialize user store") from exc
+
+    async def register(self, username: str, password: str) -> RegisteredUser:
+        password_hash = bcrypt.hashpw(
+            password.encode("utf-8"), bcrypt.gensalt()
+        ).decode("utf-8")
+        now = datetime.now(UTC)
+        user = RegisteredUser(
+            user_id=str(uuid4()),
+            api_key=secrets.token_urlsafe(32),
+            created_at=now,
+        )
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    "INSERT INTO users (user_id, username, password_hash, api_key, created_at) "
+                    "VALUES (?, ?, ?, ?, ?)",
+                    (
+                        user.user_id,
+                        username,
+                        password_hash,
+                        user.api_key,
+                        now.isoformat(),
+                    ),
+                )
+                await db.commit()
+        except aiosqlite.IntegrityError as exc:
+            raise UsernameExistsError(username) from exc
+        except aiosqlite.Error as exc:
+            raise UserStoreError("could not register user") from exc
+        return user
+
+    async def login(self, username: str, password: str) -> RegisteredUser:
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                cursor = await db.execute(
+                    "SELECT user_id, password_hash, api_key, created_at "
+                    "FROM users WHERE username = ?",
+                    (username,),
+                )
+                row = await cursor.fetchone()
+        except aiosqlite.Error as exc:
+            raise UserStoreError("could not query user") from exc
+
+        if row is None:
+            _run_constant_time_password_check()
+            raise InvalidCredentialsError("invalid username or password")
+
+        user_id, password_hash, api_key, created_at = row
+        try:
+            password_ok = bcrypt.checkpw(
+                password.encode("utf-8"), password_hash.encode("utf-8")
+            )
+        except ValueError as exc:
+            raise UserStoreError("stored password hash is corrupt") from exc
+
+        if not password_ok:
+            raise InvalidCredentialsError("invalid username or password")
+
+        try:
+            parsed_created_at = datetime.fromisoformat(created_at)
+        except ValueError as exc:
+            raise UserStoreError("stored created_at is corrupt") from exc
+
+        return RegisteredUser(
+            user_id=user_id,
+            api_key=api_key,
+            created_at=parsed_created_at,
+        )
+
+    async def authenticate(self, api_key: str | None) -> AuthenticatedUser | None:
+        if not api_key:
+            return None
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                cursor = await db.execute(
+                    "SELECT api_key, user_id FROM users WHERE api_key = ?",
+                    (api_key,),
+                )
+                row = await cursor.fetchone()
+        except aiosqlite.Error:
+            return None
+
+        if row is None:
+            return None
+        stored_key, user_id = row
+        if not hmac.compare_digest(api_key, stored_key):
+            return None
+        return AuthenticatedUser(user_id=user_id)
+
+
+def _run_constant_time_password_check() -> None:
+    """Perform a dummy bcrypt comparison to keep timing roughly constant.
+
+    Equalizes the response time between an unknown username (no hash to
+    compare against) and a wrong password for a known username, reducing the
+    feasibility of username enumeration via timing.
+    """
+    dummy_hash = bcrypt.gensalt()
+    bcrypt.checkpw(b"dummy-password", dummy_hash)