nida-tts / src /nida /data /remote_zip.py
Bilal7Dev's picture
Nida v0.5 (research preview): weights, 16 voices, code, model card
c0bce59 verified
Raw History Blame Contribute Delete
25.5 kB
"""List and extract single members of a *remote* zip archive with HTTP ``Range`` requests.
EARS ships one ~564 MB zip per speaker (ADR-0007) of which only a small part is needed. This
module reads the end-of-central-directory record (classic and ZIP64), the central directory
and then only the byte ranges of the requested members, so the archive is never downloaded
as a whole.
* :class:`RemoteZip` opens a URL. :meth:`RemoteZip.members` lists the archive,
:meth:`RemoteZip.extract` writes one member to a file (atomically: temporary file in the
destination folder, fsync, rename; a partial member never looks complete).
* Supported: stored (0) and deflate (8) members, ZIP64 end records and ZIP64 extra fields,
UTF-8 and CP437 names. Not supported (clear :class:`RemoteZipError`): encryption, other
compression methods, multi-disk archives, archives with data prepended to the zip.
* Every read goes through one *read-ahead window* (default 4 MiB): members that are close to
each other in the archive are served by a single request. Extract in ascending
``header_offset`` order to profit from it.
* Network access goes through the injectable :data:`nida.bench.fetch.Opener`; transient
failures (``URLError``, resets, short reads, HTTP 5xx/408/429) are retried with backoff of
1, 2, 4, ... seconds (capped at 30) like :func:`nida.bench.fetch.download`. A server that
answers 200 to a ranged request is refused (it would send the whole archive).
* The central directory is trusted for sizes; the local header is cross-checked (signature,
name, method). The CRC-32 and the size of every extracted member are verified, and the
decompressed size is bounded by the declared size (no zip bomb).
"""
from __future__ import annotations
import logging
import os
import re
import struct
import tempfile
import time
import urllib.error
import urllib.request
import zlib
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import IO
from nida.bench import fetch
__all__ = ["RemoteZip", "RemoteZipError", "ZipMember", "parse_central_directory"]
log = logging.getLogger(__name__)
_USER_AGENT = "nida-data/0"
_CONTENT_RANGE_RE = re.compile(r"bytes (\d+)-(\d+)/(\d+|\*)")
_MAX_BACKOFF_S = 30.0
_RETRYABLE_4XX = frozenset({408, 429})
# The transient-network set of nida.bench.fetch (same retry classification).
_NETWORK_ERRORS: tuple[type[BaseException], ...] = fetch._NETWORK_ERRORS
_COPY_STEP_BYTES = 1 << 20
_EOCD_SIG = b"PK\x05\x06"
_EOCD_FMT = "<4s4H2LH" # sig, disk, cd_disk, entries_on_disk, entries, cd_size, cd_offset, comment
_EOCD_SIZE = struct.calcsize(_EOCD_FMT)
_LOCATOR_SIG = b"PK\x06\x07"
_LOCATOR_FMT = "<4sLQL" # sig, disk_of_eocd64, eocd64_offset, total_disks
_LOCATOR_SIZE = struct.calcsize(_LOCATOR_FMT)
_EOCD64_SIG = b"PK\x06\x06"
_EOCD64_FMT = "<4sQ2H2L4Q" # sig, size, made, needed, disk, cd_disk, n_disk, n, cd_size, cd_offset
_EOCD64_SIZE = struct.calcsize(_EOCD64_FMT)
_CD_SIG = b"PK\x01\x02"
_CD_FMT = "<4s4B4HL2L5H2L"
_CD_SIZE = struct.calcsize(_CD_FMT)
_LOCAL_SIG = b"PK\x03\x04"
_LOCAL_FMT = "<4s2B4HL2L2H"
_LOCAL_SIZE = struct.calcsize(_LOCAL_FMT)
_U16 = 0xFFFF
_U32 = 0xFFFFFFFF
_ZIP64_EXTRA_ID = 0x0001
_FLAG_ENCRYPTED = 0x0001
_FLAG_STRONG_ENCRYPTION = 0x0040
_FLAG_UTF8 = 0x0800
METHOD_STORED = 0
METHOD_DEFLATED = 8
class RemoteZipError(fetch.FetchError):
"""Reading the remote archive failed (network, unsupported or corrupt zip)."""
class _Transient(Exception):
"""Internal: this request failed but an identical retry may succeed."""
@dataclass(frozen=True)
class ZipMember:
"""One central-directory entry.
Attributes:
name: Entry name (``/``-separated; a trailing ``/`` marks a directory).
raw_name: Name bytes as stored (compared with the local header).
flags: General purpose bit flags.
method: Compression method (0 stored, 8 deflate are extractable).
crc32: CRC-32 of the uncompressed data.
compressed_size: Bytes of (compressed) data.
file_size: Bytes after decompression.
header_offset: Offset of the local file header in the archive.
end_offset: Offset of the next entry (or of the central directory): everything in
``[header_offset, end_offset)`` belongs to this member.
"""
name: str
raw_name: bytes
flags: int
method: int
crc32: int
compressed_size: int
file_size: int
header_offset: int
end_offset: int
@property
def is_dir(self) -> bool:
return self.name.endswith("/")
# --------------------------------------------------------------------------------------
# Pure parsing
# --------------------------------------------------------------------------------------
@dataclass(frozen=True)
class _Directory:
"""Where the central directory is and how many entries it holds."""
offset: int
size: int
n_entries: int
def _find_eocd(tail: bytes) -> int:
"""Index of the end-of-central-directory record in ``tail`` (last valid signature).
A candidate counts only if its comment length reaches exactly the end of ``tail``, so a
signature inside the archive comment is skipped.
"""
end = len(tail)
while True:
pos = tail.rfind(_EOCD_SIG, 0, end)
if pos < 0:
raise RemoteZipError("end of central directory record not found (not a zip archive?)")
if pos + _EOCD_SIZE <= len(tail):
(comment_len,) = struct.unpack_from("<H", tail, pos + _EOCD_SIZE - 2)
if pos + _EOCD_SIZE + comment_len == len(tail):
return pos
end = pos + len(_EOCD_SIG) - 1 # exclude the occurrence at ``pos`` from the next search
def _classic_directory(eocd: bytes) -> tuple[_Directory, bool]:
"""Parse a classic EOCD; second value: True when any field is a ZIP64 sentinel."""
_, disk, cd_disk, _, n_entries, cd_size, cd_offset, _ = struct.unpack(_EOCD_FMT, eocd)
sentinel = n_entries == _U16 or cd_size == _U32 or cd_offset == _U32
if (disk, cd_disk) != (0, 0) and not sentinel:
raise RemoteZipError(f"multi-disk zip archives are not supported (disk {disk}/{cd_disk})")
return _Directory(cd_offset, cd_size, n_entries), sentinel
def _zip64_directory(record: bytes) -> _Directory:
sig, _, _, _, disk, cd_disk, _, n_entries, cd_size, cd_offset = struct.unpack(
_EOCD64_FMT, record
)
if sig != _EOCD64_SIG:
raise RemoteZipError("ZIP64 end of central directory record has a bad signature")
if (disk, cd_disk) != (0, 0):
raise RemoteZipError(f"multi-disk zip archives are not supported (disk {disk}/{cd_disk})")
return _Directory(cd_offset, cd_size, n_entries)
def _apply_zip64_extra(
extra: bytes, file_size: int, compressed_size: int, header_offset: int, name: str
) -> tuple[int, int, int]:
"""Replace sentinel (``0xFFFFFFFF``) sizes/offset by the values of the ZIP64 extra field.
The ZIP64 field holds, in this order, only those of uncompressed size, compressed size
and header offset whose 32-bit value is the sentinel.
"""
needed = [file_size == _U32, compressed_size == _U32, header_offset == _U32]
pos = 0
while pos + 4 <= len(extra):
field_id, field_size = struct.unpack_from("<HH", extra, pos)
body = extra[pos + 4 : pos + 4 + field_size]
pos += 4 + field_size
if field_id != _ZIP64_EXTRA_ID:
continue
count = sum(needed)
if len(body) < 8 * count:
raise RemoteZipError(f"{name!r}: truncated ZIP64 extra field")
values = iter(struct.unpack_from(f"<{count}Q", body, 0))
return (
next(values) if needed[0] else file_size,
next(values) if needed[1] else compressed_size,
next(values) if needed[2] else header_offset,
)
if any(needed):
raise RemoteZipError(f"{name!r}: ZIP64 sentinel without a ZIP64 extra field")
return file_size, compressed_size, header_offset
def parse_central_directory(data: bytes, n_entries: int, directory_offset: int) -> list[ZipMember]:
"""Parse ``n_entries`` central-directory entries from ``data``.
Args:
data: The central directory bytes (exactly ``cd_size`` of them).
n_entries: Entry count from the end record.
directory_offset: Archive offset of ``data``; the last member ends there.
Returns:
Members in directory order, ``end_offset`` filled in.
Raises:
RemoteZipError: on a bad signature, truncation, trailing bytes or a count mismatch.
"""
pos = 0
rows: list[tuple[str, bytes, int, int, int, int, int, int]] = []
for index in range(n_entries):
if pos + _CD_SIZE > len(data):
raise RemoteZipError(f"central directory truncated at entry {index}/{n_entries}")
fields = struct.unpack_from(_CD_FMT, data, pos)
if fields[0] != _CD_SIG:
raise RemoteZipError(f"central directory entry {index} has a bad signature")
flags, method = fields[5], fields[6]
crc, comp, size = fields[9], fields[10], fields[11]
name_len, extra_len, comment_len = fields[12], fields[13], fields[14]
offset = fields[18]
start = pos + _CD_SIZE
end = start + name_len + extra_len + comment_len
if end > len(data):
raise RemoteZipError(f"central directory truncated inside entry {index}")
raw_name = bytes(data[start : start + name_len])
extra = bytes(data[start + name_len : start + name_len + extra_len])
name = raw_name.decode("utf-8" if flags & _FLAG_UTF8 else "cp437", errors="replace")
size, comp, offset = _apply_zip64_extra(extra, size, comp, offset, name)
rows.append((name, raw_name, flags, method, crc, comp, size, offset))
pos = end
if pos != len(data):
raise RemoteZipError(f"central directory has {len(data) - pos} unparsed trailing bytes")
boundaries = sorted({row[7] for row in rows} | {directory_offset})
members: list[ZipMember] = []
for name, raw_name, flags, method, crc, comp, size, offset in rows:
later = [b for b in boundaries if b > offset]
members.append(
ZipMember(
name=name,
raw_name=raw_name,
flags=flags,
method=method,
crc32=crc,
compressed_size=comp,
file_size=size,
header_offset=offset,
end_offset=later[0] if later else directory_offset,
)
)
return members
# --------------------------------------------------------------------------------------
# RemoteZip
# --------------------------------------------------------------------------------------
class RemoteZip:
"""A zip archive behind an HTTP(S) URL that supports ``Range`` requests.
Args:
url: ``http://`` or ``https://`` URL of the archive.
opener: Network access function (injectable for tests).
timeout_s: Socket timeout per request, seconds.
max_attempts: Attempts per request (first try included), at least 1.
sleep: Sleep function used for backoff.
window_bytes: Read-ahead window; every fetch asks for at least this many bytes.
tail_bytes: Bytes read from the end of the archive to find the end record (must be
able to hold a record plus the longest comment, 65 557 bytes).
"""
def __init__(
self,
url: str,
*,
opener: fetch.Opener = fetch.default_opener,
timeout_s: float = 60.0,
max_attempts: int = 5,
sleep: Callable[[float], None] = time.sleep,
window_bytes: int = 4 << 20,
tail_bytes: int = 256 << 10,
) -> None:
if not url.lower().startswith(("http://", "https://")):
raise ValueError(f"only http(s) URLs are supported, got {url!r}")
if timeout_s <= 0:
raise ValueError(f"timeout_s must be > 0, got {timeout_s}")
if max_attempts < 1:
raise ValueError(f"max_attempts must be >= 1, got {max_attempts}")
if window_bytes < 1:
raise ValueError(f"window_bytes must be >= 1, got {window_bytes}")
if tail_bytes < _EOCD_SIZE + _LOCATOR_SIZE + 0xFFFF:
raise ValueError(f"tail_bytes must hold an end record and a comment, got {tail_bytes}")
self.url = url
self._opener = opener
self._timeout_s = timeout_s
self._max_attempts = max_attempts
self._sleep = sleep
self._window_bytes = window_bytes
self._tail_bytes = tail_bytes
self._size: int | None = None
self._members: list[ZipMember] | None = None
self._window_start = 0
self._window = b""
self.num_requests = 0
# ------------------------------------------------------------------ HTTP
def _attempt(self, start: int, stop: int) -> tuple[bytes, int | None]:
request = urllib.request.Request(
self.url, headers={"User-Agent": _USER_AGENT, "Range": f"bytes={start}-{stop - 1}"}
)
self.num_requests += 1
try:
response = self._opener(request, self._timeout_s)
except urllib.error.HTTPError as exc:
code = exc.code
exc.close()
if 400 <= code < 500 and code not in _RETRYABLE_4XX:
raise RemoteZipError(
f"HTTP {code} for {self.url} (bytes {start}-{stop - 1})"
) from exc
raise _Transient(f"HTTP {code}") from exc
except _NETWORK_ERRORS as exc:
raise _Transient(f"connecting failed: {exc!r}") from exc
with response:
if response.status == 200:
raise RemoteZipError(
f"{self.url} ignored the Range header (HTTP 200); refusing to download "
"the whole archive"
)
if response.status != 206:
raise RemoteZipError(f"unexpected HTTP status {response.status} from {self.url}")
header = response.headers.get("Content-Range")
match = _CONTENT_RANGE_RE.fullmatch(header.strip()) if header else None
if match is None:
raise RemoteZipError(
f"missing or malformed Content-Range from {self.url}: {header!r}"
)
first, last = int(match.group(1)), int(match.group(2))
total = None if match.group(3) == "*" else int(match.group(3))
if first != start or last > stop - 1 or last < first:
raise RemoteZipError(
f"{self.url}: asked for bytes {start}-{stop - 1}, server sent {first}-{last}"
)
expected = last - first + 1
chunks: list[bytes] = []
received = 0
try:
while received < expected:
chunk = response.read(min(_COPY_STEP_BYTES, expected - received))
if not chunk:
break
chunks.append(chunk)
received += len(chunk)
except _NETWORK_ERRORS as exc:
raise _Transient(f"reading failed after {received} bytes: {exc!r}") from exc
if received != expected:
raise _Transient(f"short body: {received} bytes, expected {expected}")
return b"".join(chunks), total
def _request(self, start: int, stop: int) -> tuple[bytes, int | None]:
"""Bytes ``[start, stop)`` with retries; also the archive size if the server says."""
last: _Transient | None = None
for attempt in range(1, self._max_attempts + 1):
try:
return self._attempt(start, stop)
except _Transient as exc:
last = exc
log.warning(
"range %d-%d of %s: attempt %d/%d failed: %s",
start,
stop - 1,
self.url,
attempt,
self._max_attempts,
exc,
)
if attempt < self._max_attempts:
self._sleep(min(2.0 ** (attempt - 1), _MAX_BACKOFF_S))
raise RemoteZipError(
f"range request {start}-{stop - 1} of {self.url} failed after "
f"{self._max_attempts} attempts: {last}"
) from last
@property
def size_bytes(self) -> int:
"""Archive size, learned from the ``Content-Range`` of a one-byte probe."""
if self._size is None:
_, total = self._request(0, 1)
if total is None or total < _EOCD_SIZE:
raise RemoteZipError(f"cannot determine the size of {self.url} (got {total!r})")
self._size = total
return self._size
def _read(self, start: int, length: int) -> bytes:
"""``length`` bytes at ``start``, served from the read-ahead window when possible."""
size = self.size_bytes
if length < 0 or start < 0 or start + length > size:
raise RemoteZipError(
f"{self.url}: read of {length} bytes at {start} is outside the archive ({size})"
)
if length == 0:
return b""
offset = start - self._window_start
if offset >= 0 and offset + length <= len(self._window):
return self._window[offset : offset + length]
stop = min(size, start + max(length, self._window_bytes))
body, _ = self._request(start, stop)
if len(body) != stop - start:
raise RemoteZipError(
f"{self.url}: server returned {len(body)} bytes for a {stop - start}-byte range"
)
self._window_start, self._window = start, body
return body[:length]
# ------------------------------------------------------------------ listing
def _locate_directory(self) -> tuple[_Directory, int]:
"""Central directory location and the archive offset where its end record begins."""
size = self.size_bytes
tail_len = min(size, self._tail_bytes)
tail = self._read(size - tail_len, tail_len)
eocd_at = _find_eocd(tail)
eocd_pos = size - tail_len + eocd_at
directory, sentinel = _classic_directory(tail[eocd_at : eocd_at + _EOCD_SIZE])
locator_at = eocd_at - _LOCATOR_SIZE
has_locator = locator_at >= 0 and tail[locator_at : locator_at + 4] == _LOCATOR_SIG
if not has_locator:
if sentinel:
raise RemoteZipError("ZIP64 sentinel values but no ZIP64 end locator")
return directory, eocd_pos
_, locator_disk, eocd64_offset, total_disks = struct.unpack_from(
_LOCATOR_FMT, tail, locator_at
)
if locator_disk != 0 or total_disks > 1:
raise RemoteZipError("multi-disk zip archives are not supported")
if eocd64_offset + _EOCD64_SIZE > size:
raise RemoteZipError("ZIP64 end record offset lies outside the archive")
directory = _zip64_directory(self._read(eocd64_offset, _EOCD64_SIZE))
return directory, eocd64_offset
def members(self) -> list[ZipMember]:
"""All central-directory entries (directories included), in directory order.
The result is cached for the lifetime of the object.
Raises:
RemoteZipError: on network failure or an unsupported or corrupt archive.
"""
if self._members is None:
directory, end_limit = self._locate_directory()
if directory.offset + directory.size > end_limit:
raise RemoteZipError(
f"central directory ({directory.offset}+{directory.size}) overlaps the end "
f"record at {end_limit}; archives with prepended data are not supported"
)
data = self._read(directory.offset, directory.size)
self._members = parse_central_directory(data, directory.n_entries, directory.offset)
log.info("%s: %d entries, %d bytes", self.url, len(self._members), self.size_bytes)
return list(self._members)
# ------------------------------------------------------------------ extraction
def extract(self, member: ZipMember, dest: Path, *, skip_existing: bool = True) -> Path:
"""Write one member to ``dest`` (atomic) and verify its size and CRC-32.
Args:
member: An entry of :meth:`members`.
dest: Destination file; parent folders are created.
skip_existing: Return ``dest`` without any network access when it already has the
member's size (files only ever appear complete, see the module docstring).
Returns:
``dest``.
Raises:
RemoteZipError: for a directory, an encrypted or unsupported member, a local
header that contradicts the central directory, corrupt data, or network
failure.
"""
if member.is_dir:
raise RemoteZipError(f"{member.name!r} is a directory")
if member.flags & (_FLAG_ENCRYPTED | _FLAG_STRONG_ENCRYPTION):
raise RemoteZipError(f"{member.name!r} is encrypted")
if member.method not in (METHOD_STORED, METHOD_DEFLATED):
raise RemoteZipError(f"{member.name!r}: unsupported compression method {member.method}")
if skip_existing and dest.is_file() and dest.stat().st_size == member.file_size:
return dest
data_start = self._data_start(member)
dest.parent.mkdir(parents=True, exist_ok=True)
fd, tmp_name = tempfile.mkstemp(prefix=f".{dest.name}.", suffix=".tmp", dir=dest.parent)
tmp_path = Path(tmp_name)
try:
with os.fdopen(fd, "wb") as handle:
self._copy_data(member, data_start, handle)
handle.flush()
os.fsync(handle.fileno())
tmp_path.replace(dest)
except BaseException:
tmp_path.unlink(missing_ok=True)
raise
return dest
def _data_start(self, member: ZipMember) -> int:
"""Archive offset of the member's data, after checking its local header."""
header = self._read(member.header_offset, _LOCAL_SIZE)
fields = struct.unpack(_LOCAL_FMT, header)
if fields[0] != _LOCAL_SIG:
raise RemoteZipError(f"{member.name!r}: bad local header signature")
method, name_len, extra_len = int(fields[4]), int(fields[10]), int(fields[11])
if method != member.method:
raise RemoteZipError(
f"{member.name!r}: local header method {method} != central directory "
f"{member.method}"
)
local_name = self._read(member.header_offset + _LOCAL_SIZE, name_len)
if local_name != member.raw_name:
raise RemoteZipError(f"{member.name!r}: local header names {local_name!r}")
data_start = member.header_offset + _LOCAL_SIZE + name_len + extra_len
if data_start + member.compressed_size > member.end_offset:
raise RemoteZipError(
f"{member.name!r}: data ({data_start}+{member.compressed_size}) runs into the "
f"next entry at {member.end_offset}"
)
return data_start
def _copy_data(self, member: ZipMember, data_start: int, handle: IO[bytes]) -> None:
write = handle.write
inflater = zlib.decompressobj(-zlib.MAX_WBITS) if member.method == METHOD_DEFLATED else None
crc = 0
written = 0
position = data_start
remaining = member.compressed_size
while remaining > 0:
chunk = self._read(position, min(remaining, self._window_bytes))
position += len(chunk)
remaining -= len(chunk)
pending = chunk
while pending:
if inflater is None:
out, pending = pending, b""
else:
try:
out = inflater.decompress(pending, _COPY_STEP_BYTES)
except zlib.error as exc:
raise RemoteZipError(
f"{member.name!r}: corrupt deflate data: {exc}"
) from exc
pending = inflater.unconsumed_tail
written += len(out)
if written > member.file_size:
raise RemoteZipError(
f"{member.name!r}: more than the declared {member.file_size} bytes"
)
crc = zlib.crc32(out, crc)
write(out)
if inflater is not None:
try:
tail = inflater.flush()
except zlib.error as exc:
raise RemoteZipError(f"{member.name!r}: corrupt deflate data: {exc}") from exc
written += len(tail)
crc = zlib.crc32(tail, crc)
write(tail)
if not inflater.eof:
raise RemoteZipError(f"{member.name!r}: deflate stream is truncated")
if written != member.file_size:
raise RemoteZipError(
f"{member.name!r}: {written} bytes extracted, expected {member.file_size}"
)
if crc != member.crc32:
raise RemoteZipError(
f"{member.name!r}: CRC-32 {crc:08x} does not match the archive's {member.crc32:08x}"
)