A02/A03: Add state validation and normalize subdomain keys for shared cooldown

This commit is contained in:
ik
2026-09-10 20:34:05 +07:00
parent f29ec706fd
commit 57aa3ffafb
2 changed files with 94 additions and 10 deletions
+38 -5
View File
@@ -44,12 +44,25 @@ MAX_RESPONSE_BYTES = 5 * 1024 * 1024 # 5 MB
def _read_state(state_file: Path) -> dict:
"""Read state file with shared lock; return empty dict if missing/corrupt."""
"""Read state file with shared lock; return empty dict if missing/corrupt.
Validates that the parsed JSON is a flat dict of string->number pairs.
Returns {} for missing, corrupt, or structurally invalid files.
"""
try:
with open(state_file, "r") as f:
fcntl.flock(f, fcntl.LOCK_SH)
try:
return json.loads(f.read())
data = json.loads(f.read())
if not isinstance(data, dict):
return {}
# Validate structure: flat dict of string->number
for key, value in data.items():
if not isinstance(key, str):
return {}
if not isinstance(value, (int, float)):
return {}
return data
finally:
fcntl.flock(f, fcntl.LOCK_UN)
except (FileNotFoundError, json.JSONDecodeError, ValueError, OSError):
@@ -76,10 +89,21 @@ def _write_state(state_file: Path, state: dict) -> None:
def fetch_site_key(url: str) -> str:
"""Return a stable cooldown key shared by all endpoints of one site."""
"""Return a stable cooldown key shared by all endpoints of one site.
Normalizes hostname to a common key for related domains:
- Strips common subdomains (www, download, api, cdn)
- Keeps the base domain as the key
"""
hostname = (urlsplit(url).hostname or "").lower()
if hostname.startswith("www."):
hostname = hostname[4:]
if hostname.startswith("download."):
hostname = hostname[9:]
if hostname.startswith("api."):
hostname = hostname[4:]
if hostname.startswith("cdn."):
hostname = hostname[4:]
if not hostname:
raise ValueError("source URL must include a hostname")
return hostname
@@ -105,6 +129,8 @@ def check_and_reserve(
Opens state file with exclusive lock, reads state, checks cooldown,
reserves if allowed — all in one critical section. Uses lockfile
pattern for cross-process coordination.
Handles corrupt/missing state files gracefully by treating them as empty.
"""
if now is None:
now = time.time()
@@ -115,10 +141,17 @@ def check_and_reserve(
with open(lock_file, "w") as lf:
fcntl.flock(lf, fcntl.LOCK_EX)
try:
# Read state under lock
# Read state under lock; validate structure
try:
with open(state_file, "r") as sf:
state = json.loads(sf.read()) or {}
raw = sf.read()
state = json.loads(raw) or {}
if not isinstance(state, dict):
state = {}
for key, value in state.items():
if not isinstance(key, str) or not isinstance(value, (int, float)):
state = {}
break
except (FileNotFoundError, json.JSONDecodeError, ValueError, OSError):
state = {}
# Check cooldown