mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 20:15:10 +00:00
imSp4rky
d4690e5c6b
Request 4 MiB sub-ranges once a host cuts a body, skip the session's own urllib3 Retry inside download(), fail the ranged path at once on a 4xx, and cap resumes at 50 of 1 MiB each. Adds a deny fault profile.
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
import sys
|
|
import threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
import pytest
|
|
import requests as rq
|
|
|
|
from unshackle.core.constants import DOWNLOAD_CANCELLED
|
|
from unshackle.core.downloaders.requests import download
|
|
|
|
# the package re-exports the `requests` function, shadowing the module of the same name
|
|
downloader = sys.modules[download.__module__]
|
|
|
|
BODY = bytes(range(256)) * 256 # 64 KiB
|
|
CUT_AFTER = 10 * 1024
|
|
CHUNK = 16 * 1024
|
|
|
|
|
|
class _CutServer(ThreadingHTTPServer):
|
|
"""Serves BODY but ends any body longer than CUT_AFTER early, like a per-request-limited CDN."""
|
|
|
|
allow_reuse_address = True
|
|
daemon_threads = True
|
|
|
|
def __init__(self):
|
|
self.ranges: list[str] = []
|
|
self.cuts = 0
|
|
self.connections = 0
|
|
super().__init__(("127.0.0.1", 0), _CutHandler)
|
|
self.url = f"http://127.0.0.1:{self.server_address[1]}/seg.mp4"
|
|
|
|
def handle_error(self, request, client_address): # the client hanging up is expected here
|
|
pass
|
|
|
|
|
|
class _CutHandler(BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
def setup(self):
|
|
super().setup()
|
|
self.server.connections += 1
|
|
|
|
def do_GET(self):
|
|
header = self.headers.get("Range")
|
|
self.server.ranges.append(header)
|
|
start, end = 0, len(BODY) - 1
|
|
if header:
|
|
start_s, _, end_s = header.removeprefix("bytes=").partition("-")
|
|
start, end = int(start_s), (int(end_s) if end_s else len(BODY) - 1)
|
|
body = BODY[start : end + 1]
|
|
self.send_response(206 if header else 200)
|
|
if header:
|
|
self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
if len(body) > CUT_AFTER:
|
|
self.server.cuts += 1
|
|
self.wfile.write(body[:CUT_AFTER])
|
|
self.close_connection = True
|
|
return
|
|
self.wfile.write(body)
|
|
|
|
|
|
@pytest.fixture
|
|
def server(monkeypatch):
|
|
DOWNLOAD_CANCELLED.clear()
|
|
monkeypatch.setattr(downloader, "RETRY_WAIT", 0)
|
|
monkeypatch.setattr(downloader, "REQUEST_CHUNK_SIZE", CHUNK)
|
|
monkeypatch.setattr(downloader, "MIN_REQUEST_CHUNK", 4 * 1024)
|
|
monkeypatch.setattr(downloader, "MIN_RESUME_PROGRESS", 1024)
|
|
server = _CutServer()
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
yield server
|
|
server.shutdown()
|
|
server.server_close()
|
|
DOWNLOAD_CANCELLED.clear()
|
|
|
|
|
|
def test_a_cut_body_bounds_the_following_requests(tmp_path, server):
|
|
save_path = tmp_path / "0.mp4"
|
|
list(download(url=server.url, save_path=save_path, session=rq.Session()))
|
|
|
|
assert save_path.read_bytes() == BODY
|
|
# the open-ended first request is cut, the chunk-sized one after it is cut too, and the
|
|
# halved chunk fits under the host's limit, so every later request completes whole
|
|
assert server.cuts == 2
|
|
spans = [r.removeprefix("bytes=").split("-") for r in server.ranges if r and r.split("-")[-1]]
|
|
assert spans and all(int(end) - int(start) + 1 <= CHUNK for start, end in spans)
|
|
# the bounded requests are cheap because they share one pooled connection
|
|
assert server.connections < len(server.ranges)
|