fix: resolve concurrency and DoS pitfalls from PR #272 review

- index/utils.py: fix an asyncio deadlock in _sync_llm_semaphore. Sync LLM
  calls (check_toc / process_no_toc / toc_transformer) run on the event-loop
  thread nested inside the async meta_processor, and its blocking ceiling
  acquire could wait forever for a permit held by async _llm_semaphore holders
  that can only release it once the (now-frozen) loop runs. Take the slot
  non-blocking when on a running loop; keep the blocking acquire off-loop.
- index/utils.py: bound parse_pages ranges before materializing range() into
  the list, so a huge span like '1-2000000000' is rejected up front instead of
  exhausting memory before the 1000-page cap is ever checked (DoS).
- storage/sqlite.py: bump a generation counter on close() so a thread that
  cached a connection in thread-local storage reconnects on its next call
  instead of reusing a closed handle (ProgrammingError).
- backend/cloud.py: after connect, bail out of the SSE background thread if the
  consumer already abandoned the stream, instead of draining it in the background.

Adds regression tests for each fix.
This commit is contained in:
mountain 2026-07-10 14:49:55 +08:00
parent d97231b480
commit 9ad54122bb
6 changed files with 192 additions and 10 deletions

View file

@ -139,6 +139,46 @@ def test_close_closes_connections_created_in_other_threads(storage):
conns["worker"].execute("SELECT 1")
def test_worker_reconnects_via_get_conn_after_close(storage):
"""Regression: after close(), a thread that had already cached a connection
in thread-local storage would get that now-CLOSED handle back from
_get_conn (close() can only del its own thread-local), raising
ProgrammingError instead of transparently reconnecting. A generation bump
on close() must make the SAME thread's next _get_conn hand back a fresh,
working connection."""
import threading
storage.create_collection("papers")
cached = threading.Event()
closed = threading.Event()
result = {}
def worker():
# 1. cache a connection in this thread's thread-local
storage._get_conn().execute("SELECT 1")
cached.set()
# 2. wait until the main thread closed the storage (invalidating it)
closed.wait(timeout=5)
# 3. reuse from the SAME thread -> must reconnect, not reuse closed conn
try:
result["val"] = storage._get_conn().execute("SELECT 1").fetchone()[0]
result["list"] = storage.list_collections()
except Exception as e: # noqa: BLE001 - record for assertion
result["err"] = f"{type(e).__name__}: {e}"
t = threading.Thread(target=worker)
t.start()
cached.wait(timeout=5)
storage.close() # closes + invalidates the worker's cached connection
closed.set()
t.join(timeout=5)
assert "err" not in result, f"reconnect after close failed: {result.get('err')}"
assert result["val"] == 1
assert result["list"] == ["papers"]
def test_duplicate_file_hash_in_collection_raises(storage):
"""UNIQUE(collection_name, file_hash) guards the add-same-file race."""
import sqlite3