mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/sources/catalog/channel_client.py

@@ -0,0 +1,296 @@
+"""Fetching a pinned catalogue generation off the release channel, and installing it.
+
+Two fetches and no discovery. The pin names the manifest's address and the manifest's exact
+digest; the manifest names the archive's object name and the archive's exact digest; the archive
+is then handed to ``update.install_release``, which is the same code path a hand-typed
+``mr-data catalog-update`` goes through and performs the same real store verification. Nothing
+here decides *which* generation to install -- that decision is the pin, and the pin is
+deploy-owned.
+
+The transport is the standard library's, hardened the way every other one in this repository is:
+ambient proxies off, redirects off, a byte ceiling per response, a timeout. Redirects are off
+because the archive URL is derived from the pinned channel address, and a redirect would let the
+channel host move the fetch somewhere the pin never named. The digest would still catch it, but
+"the digest would still catch it" is not a reason to follow it.
+
+AUTHORISATION. The release bucket has public access prevention enforced, so a reader presents a
+bearer token. On a worker that token comes from the instance metadata server and is scoped to the
+worker's own service account; nothing about it is ever written down, logged, or put in a URL. A
+channel that is world-readable needs no token and this module sends none.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import ssl
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, Protocol
+from urllib.error import HTTPError
+from urllib.parse import urlsplit, urlunsplit
+from urllib.request import HTTPRedirectHandler, HTTPSHandler, ProxyHandler, Request, build_opener
+
+from mostlyright.data_harness.sources.catalog.channel import (
+    CHANNEL_MANIFEST_OBJECT,
+    MAX_CHANNEL_MANIFEST_BYTES,
+    CatalogChannelPin,
+    ChannelManifest,
+    verify_channel_manifest,
+)
+from mostlyright.data_harness.sources.catalog.update import (
+    MAX_RELEASE_ARCHIVE_BYTES,
+    CatalogUpdateReport,
+    confirm_already_current,
+    install_release,
+)
+from mostlyright.data_harness.sources.contracts import SourceContractError
+
+__all__ = [
+    "CHANNEL_TIMEOUT_SECONDS",
+    "METADATA_TOKEN_URL",
+    "ChannelTransport",
+    "UrlLibChannelTransport",
+    "archive_url",
+    "fetch_channel_manifest",
+    "fetch_release_archive",
+    "install_from_channel",
+    "metadata_access_token",
+]
+
+#: An archive is up to a gibibyte, so the read budget is the one ``catalog_update`` already gives a
+#: release download rather than the thirty seconds an ordinary JSON call gets.
+CHANNEL_TIMEOUT_SECONDS = 600.0
+
+#: Where a Google-hosted worker asks for its own service account's access token.
+METADATA_TOKEN_URL = (
+    "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token"
+)
+
+MAX_ACCESS_TOKEN_RESPONSE_BYTES = 16 * 1024
+
+
+class ChannelTransport(Protocol):
+    """The one seam the tests substitute: status, body bytes, lowercased response headers."""
+
+    def request(
+        self,
+        method: str,
+        url: str,
+        headers: Mapping[str, str],
+        body: bytes | None,
+        maximum: int,
+    ) -> tuple[int, bytes, Mapping[str, str]]: ...
+
+
+class _NoRedirect(HTTPRedirectHandler):
+    """Refuse every redirect. The pin named an address; a redirect is a different address."""
+
+    def redirect_request(
+        self,
+        request: Request,
+        file_pointer: Any,
+        code: int,
+        message: str,
+        headers: Any,
+        new_url: str,
+    ) -> None:
+        return None
+
+
+class UrlLibChannelTransport:
+    """Bounded HTTPS with ambient proxies and redirects disabled."""
+
+    def __init__(self, *, timeout: float = CHANNEL_TIMEOUT_SECONDS) -> None:
+        self._timeout = timeout
+        self._opener = build_opener(
+            ProxyHandler({}), HTTPSHandler(context=ssl.create_default_context()), _NoRedirect()
+        )
+
+    def request(
+        self,
+        method: str,
+        url: str,
+        headers: Mapping[str, str],
+        body: bytes | None,
+        maximum: int,
+    ) -> tuple[int, bytes, Mapping[str, str]]:
+        request = Request(url, data=body, headers=dict(headers), method=method)
+        try:
+            response = self._opener.open(request, timeout=self._timeout)
+        except HTTPError as error:
+            response = error
+        with response:
+            raw = response.read(maximum + 1)
+            status = int(response.status)
+            response_headers = {key.lower(): value for key, value in response.headers.items()}
+        if len(raw) > maximum:
+            raise SourceContractError(
+                "CATALOG_CHANNEL_FETCH", url, "the response exceeded the byte ceiling for its route"
+            )
+        return status, raw, response_headers
+
+
+def _get(
+    transport: ChannelTransport,
+    url: str,
+    *,
+    maximum: int,
+    authorization: str | None,
+    label: str,
+) -> bytes:
+    headers = {"Accept": "*/*"}
+    if authorization is not None:
+        headers["Authorization"] = authorization
+    status, raw, _headers = transport.request("GET", url, headers, None, maximum)
+    if status != 200:
+        # The status is a fact about the channel, not about the caller's credential, so it is
+        # reported. The credential itself is never in this message because it is never in `url`.
+        raise SourceContractError(
+            "CATALOG_CHANNEL_FETCH", url, f"the channel answered {status} for the {label}"
+        )
+    return raw
+
+
+def archive_url(channel_url: str, manifest: ChannelManifest) -> str:
+    """The address of the archive this manifest names: the channel's own folder, one object over.
+
+    Derived rather than published so a manifest cannot point a reader at a different host than the
+    one the pin authorised. The manifest chooses the object name; the pin chooses everywhere else.
+    """
+
+    parsed = urlsplit(channel_url)
+    prefix = parsed.path[: -len(CHANNEL_MANIFEST_OBJECT)]
+    return urlunsplit(
+        (parsed.scheme, parsed.netloc, prefix + manifest.release.archive_object, "", "")
+    )
+
+
+def fetch_channel_manifest(
+    pin: CatalogChannelPin,
+    *,
+    transport: ChannelTransport | None = None,
+    authorization: str | None = None,
+) -> ChannelManifest:
+    """Fetch the channel manifest and believe it only if it is the pinned bytes."""
+
+    selected = transport or UrlLibChannelTransport()
+    raw = _get(
+        selected,
+        pin.channel_url,
+        maximum=MAX_CHANNEL_MANIFEST_BYTES,
+        authorization=authorization,
+        label="channel manifest",
+    )
+    return verify_channel_manifest(raw, expected_manifest_sha256=pin.manifest_sha256)
+
+
+def fetch_release_archive(
+    pin: CatalogChannelPin,
+    manifest: ChannelManifest,
+    *,
+    transport: ChannelTransport | None = None,
+    authorization: str | None = None,
+) -> bytes:
+    """Fetch the archive the manifest names, and refuse the bytes unless they are the bytes."""
+
+    selected = transport or UrlLibChannelTransport()
+    url = archive_url(pin.channel_url, manifest)
+    raw = _get(
+        selected,
+        url,
+        maximum=min(manifest.release.archive_bytes, MAX_RELEASE_ARCHIVE_BYTES),
+        authorization=authorization,
+        label="release archive",
+    )
+    if len(raw) != manifest.release.archive_bytes:
+        raise SourceContractError(
+            "CATALOG_CHANNEL_ARCHIVE",
+            url,
+            "the archive is not the byte count the channel manifest states",
+        )
+    if hashlib.sha256(raw).hexdigest() != manifest.release.archive_sha256:
+        raise SourceContractError(
+            "CATALOG_CHANNEL_ARCHIVE",
+            url,
+            "the archive bytes do not match the digest the channel manifest states",
+        )
+    return raw
+
+
+def metadata_access_token(
+    *,
+    transport: ChannelTransport | None = None,
+    scope: str = "https://www.googleapis.com/auth/devstorage.read_only",
+) -> str:
+    """This worker's own access token, as one ``Authorization`` header value.
+
+    Read from the instance metadata server, which is the only place a hosted worker has an
+    ambient identity, and asked for the narrowest scope that can read one object.
+    """
+
+    selected = transport or UrlLibChannelTransport(timeout=10.0)
+    from urllib.parse import urlencode
+
+    url = METADATA_TOKEN_URL + "?" + urlencode({"scopes": scope})
+    status, raw, _headers = selected.request(
+        "GET",
+        url,
+        {"Metadata-Flavor": "Google"},
+        None,
+        MAX_ACCESS_TOKEN_RESPONSE_BYTES,
+    )
+    if status != 200:
+        raise SourceContractError(
+            "CATALOG_CHANNEL_IDENTITY",
+            "metadata",
+            f"the metadata server answered {status} for this worker's access token",
+        )
+    from mostlyright.data_harness.canonical import CanonicalJSONError, parse_json
+
+    try:
+        payload = parse_json(raw)
+    except CanonicalJSONError as error:
+        raise SourceContractError(
+            "CATALOG_CHANNEL_IDENTITY", "metadata", "the access token response is not JSON"
+        ) from error
+    token = payload.get("access_token") if isinstance(payload, dict) else None
+    if not isinstance(token, str) or not token or len(token) > 8192:
+        raise SourceContractError(
+            "CATALOG_CHANNEL_IDENTITY", "metadata", "the access token response carried no token"
+        )
+    return "Bearer " + token
+
+
+def install_from_channel(
+    catalog_root: Path,
+    pin: CatalogChannelPin,
+    *,
+    transport: ChannelTransport | None = None,
+    authorization: str | None = None,
+    replace: bool = False,
+    neural_backend: Any = None,
+) -> CatalogUpdateReport:
+    """Install the generation this pin names, fetching only what is not already installed.
+
+    The manifest is fetched first even when the generation turns out to be present: the pin names
+    a manifest digest, not an archive digest, so there is no way to know which archive is meant
+    without reading the manifest. Once it is read, an already-installed generation short-circuits
+    the archive download entirely -- which is what makes fetch-at-start cheap on a warm worker.
+    """
+
+    manifest = fetch_channel_manifest(pin, transport=transport, authorization=authorization)
+    current = confirm_already_current(
+        catalog_root,
+        expected_archive_sha256=manifest.release.archive_sha256,
+        neural_backend=neural_backend,
+    )
+    if current is not None:
+        return current
+    raw = fetch_release_archive(pin, manifest, transport=transport, authorization=authorization)
+    return install_release(
+        catalog_root,
+        raw,
+        expected_archive_sha256=manifest.release.archive_sha256,
+        replace=replace,
+        neural_backend=neural_backend,
+    )