Spaces:
Running
Running
Isolate DNS pins per concurrent download
Browse files
app/services/downloader.py
CHANGED
|
@@ -23,12 +23,14 @@ class _PinnedNetworkBackend:
|
|
| 23 |
preserving SNI and certificate hostname verification.
|
| 24 |
"""
|
| 25 |
|
| 26 |
-
def __init__(self,
|
| 27 |
-
|
|
|
|
|
|
|
| 28 |
self.fallback = fallback
|
| 29 |
|
| 30 |
async def connect_tcp(self, host: str, port: int, timeout: float | None = None, local_address: str | None = None, socket_options=None):
|
| 31 |
-
pinned = self.
|
| 32 |
targets = pinned if pinned is not None else (host,)
|
| 33 |
last_error: Exception | None = None
|
| 34 |
for target in targets:
|
|
@@ -59,7 +61,6 @@ class Downloader:
|
|
| 59 |
def __init__(self, settings: Settings, validator: MediaValidator) -> None:
|
| 60 |
self.settings = settings
|
| 61 |
self.validator = validator
|
| 62 |
-
self._pinned: dict[tuple[str, int], tuple[str, ...]] = {}
|
| 63 |
|
| 64 |
async def download(self, url: str, destination_dir: Path, filename: str | None = None) -> tuple[Path, str]:
|
| 65 |
current_url = url
|
|
@@ -69,9 +70,10 @@ class Downloader:
|
|
| 69 |
# HTTPX 0.28 uses httpcore's network backend for socket creation. The
|
| 70 |
# wrapper pins the TCP destination while httpcore retains the URL host
|
| 71 |
# for TLS SNI, certificate checks, and the HTTP Host header.
|
|
|
|
| 72 |
try:
|
| 73 |
from httpcore._backends.auto import AutoBackend
|
| 74 |
-
transport._pool._network_backend = _PinnedNetworkBackend(
|
| 75 |
except (AttributeError, ImportError) as exc:
|
| 76 |
raise DownloadError("Secure connection-pinned downloads are unavailable in this runtime") from exc
|
| 77 |
|
|
@@ -80,7 +82,7 @@ class Downloader:
|
|
| 80 |
trust_env=False, transport=transport,
|
| 81 |
) as client:
|
| 82 |
for _ in range(6):
|
| 83 |
-
await self.validate_url(current_url)
|
| 84 |
destination: Path | None = None
|
| 85 |
try:
|
| 86 |
async with client.stream("GET", current_url, follow_redirects=False) as response:
|
|
|
|
| 23 |
preserving SNI and certificate hostname verification.
|
| 24 |
"""
|
| 25 |
|
| 26 |
+
def __init__(self, pinned: dict[tuple[str, int], tuple[str, ...]], fallback: object) -> None:
|
| 27 |
+
# Per-download mapping prevents concurrent requests from overwriting one
|
| 28 |
+
# another's validated DNS results for the same hostname.
|
| 29 |
+
self.pinned = pinned
|
| 30 |
self.fallback = fallback
|
| 31 |
|
| 32 |
async def connect_tcp(self, host: str, port: int, timeout: float | None = None, local_address: str | None = None, socket_options=None):
|
| 33 |
+
pinned = self.pinned.get((host.lower().rstrip("."), port))
|
| 34 |
targets = pinned if pinned is not None else (host,)
|
| 35 |
last_error: Exception | None = None
|
| 36 |
for target in targets:
|
|
|
|
| 61 |
def __init__(self, settings: Settings, validator: MediaValidator) -> None:
|
| 62 |
self.settings = settings
|
| 63 |
self.validator = validator
|
|
|
|
| 64 |
|
| 65 |
async def download(self, url: str, destination_dir: Path, filename: str | None = None) -> tuple[Path, str]:
|
| 66 |
current_url = url
|
|
|
|
| 70 |
# HTTPX 0.28 uses httpcore's network backend for socket creation. The
|
| 71 |
# wrapper pins the TCP destination while httpcore retains the URL host
|
| 72 |
# for TLS SNI, certificate checks, and the HTTP Host header.
|
| 73 |
+
pinned: dict[tuple[str, int], tuple[str, ...]] = {}
|
| 74 |
try:
|
| 75 |
from httpcore._backends.auto import AutoBackend
|
| 76 |
+
transport._pool._network_backend = _PinnedNetworkBackend(pinned, AutoBackend())
|
| 77 |
except (AttributeError, ImportError) as exc:
|
| 78 |
raise DownloadError("Secure connection-pinned downloads are unavailable in this runtime") from exc
|
| 79 |
|
|
|
|
| 82 |
trust_env=False, transport=transport,
|
| 83 |
) as client:
|
| 84 |
for _ in range(6):
|
| 85 |
+
await self.validate_url(current_url, pinned=pinned)
|
| 86 |
destination: Path | None = None
|
| 87 |
try:
|
| 88 |
async with client.stream("GET", current_url, follow_redirects=False) as response:
|