fix auth sqlite connection leak
This commit is contained in:
parent
55c8415f75
commit
dafd78b71b
34
core/auth.py
34
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
|
||||
|
||||
# 复用外部事务时只执行删除,不在这里额外提交。
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user