diff --git a/core/auth.py b/core/auth.py index c0f0155..a3073a1 100644 --- a/core/auth.py +++ b/core/auth.py @@ -6,9 +6,11 @@ import secrets import sqlite3 import threading import time +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timezone from http.cookies import SimpleCookie +from typing import Iterator from .config import ( AUTH_COOKIE_SECURE, @@ -65,7 +67,7 @@ def ensure_auth_store() -> None: # 先确保数据库目录存在,SQLite 才能正常创建文件。 AUTH_DB_PATH.parent.mkdir(parents=True, exist_ok=True) - with connect_db() as connection: + with db_connection() as connection: # 单机线程服务读多写少,WAL 更适合并发读场景。 connection.execute("PRAGMA journal_mode = WAL") # 用户表只保存账号和密码摘要。 @@ -122,8 +124,6 @@ def ensure_auth_store() -> None: ), ) - connection.commit() - # 只有建表和默认账号都准备好后,才把进程状态切成已初始化。 INITIALIZED = True @@ -136,6 +136,20 @@ def connect_db() -> sqlite3.Connection: return connection +@contextmanager +def db_connection() -> Iterator[sqlite3.Connection]: + # sqlite3.Connection 的 with 语法不会关闭连接,统一在这里提交并关闭。 + connection = connect_db() + try: + yield connection + connection.commit() + except Exception: + connection.rollback() + raise + finally: + connection.close() + + def hash_password(password: str, *, salt: str | None = None) -> str: # 没传 salt 时生成新的随机盐值。 actual_salt = salt or secrets.token_hex(16) @@ -184,7 +198,7 @@ def authenticate_user(username: str, password: str) -> str | None: if not username or not password: return None - with connect_db() as connection: + with db_connection() as connection: # 当前账号只支持精确匹配,避免大小写模糊带来歧义。 row = connection.execute( "SELECT username, password_hash FROM users WHERE username = ?", @@ -214,7 +228,7 @@ def create_session(username: str, *, remote_addr: str = "", user_agent: str = "" expires_at_ts = int(time.time()) + max(1, AUTH_SESSION_TTL_HOURS) * 3600 now = utc_now() - with connect_db() as connection: + with db_connection() as connection: # 每次登录前顺手清理一遍过期会话,避免本地库无限增长。 purge_expired_sessions(connection=connection) # 新会话和来源信息一起落库,后续排查时能知道来源。 @@ -230,7 +244,6 @@ def create_session(username: str, *, remote_addr: str = "", user_agent: str = "" "UPDATE users SET last_login_at = ? WHERE username = ?", (now, username), ) - connection.commit() # 把可直接写回前端的会话信息一起返回。 return CreatedSession( @@ -251,7 +264,7 @@ def get_session(token: str) -> SessionRecord | None: now_ts = int(time.time()) token_hash = hash_token(token) - with connect_db() as connection: + with db_connection() as connection: # 先删掉已过期的历史会话,减少后续查询噪音。 purge_expired_sessions(connection=connection, now_ts=now_ts) row = connection.execute( @@ -262,7 +275,6 @@ def get_session(token: str) -> SessionRecord | None: """, (token_hash, now_ts), ).fetchone() - connection.commit() # 没查到时说明 token 不存在或已经过期。 if row is None: @@ -283,10 +295,9 @@ def delete_session(token: str) -> None: if not token: return - with connect_db() as connection: + with db_connection() as connection: # 删除不存在的 token 也视为成功,便于前端幂等退出。 connection.execute("DELETE FROM sessions WHERE token_hash = ?", (hash_token(token),)) - connection.commit() def purge_expired_sessions(*, connection: sqlite3.Connection | None = None, now_ts: int | None = None) -> None: @@ -295,9 +306,8 @@ def purge_expired_sessions(*, connection: sqlite3.Connection | None = None, now_ # 外部没传连接时自行开一个短连接完成清理。 if connection is None: - with connect_db() as owned_connection: + with db_connection() as owned_connection: owned_connection.execute("DELETE FROM sessions WHERE expires_at <= ?", (actual_now_ts,)) - owned_connection.commit() return # 复用外部事务时只执行删除,不在这里额外提交。