diff --git a/.github/workflows/conventional-commits.yml b/.github/workflows/conventional-commits.yml deleted file mode 100644 index f349c50c..00000000 --- a/.github/workflows/conventional-commits.yml +++ /dev/null @@ -1,22 +0,0 @@ -name: PR Conventional Commit Validation - -on: - pull_request_target: - types: [opened, synchronize, reopened, edited] - -jobs: - validate-pr-title: - runs-on: ubuntu-latest - steps: - - uses: actions/create-github-app-token@v3.2.0 - id: app-token - with: - app-id: ${{ secrets.DS_RELEASE_BOT_ID }} - private-key: ${{ secrets.DS_RELEASE_BOT_PRIVATE_KEY }} - permission-pull-requests: write - - - name: PR Conventional Commit Validation - uses: ytanikin/pr-conventional-commits@1.5.2 - with: - task_types: '["feat","fix","docs","test","ci","refactor","perf","chore","revert"]' - token: ${{ steps.app-token.outputs.token }} diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml index e87f75c1..402d10a5 100644 --- a/.github/workflows/wheels.yml +++ b/.github/workflows/wheels.yml @@ -223,7 +223,7 @@ jobs: name: Release environment: name: pypi-release - url: https://pypi.org/p/obstore + url: https://pypi.org/p/obstore-databaas permissions: # IMPORTANT: this permission is mandatory for trusted publishing id-token: write diff --git a/CHANGELOG.md b/CHANGELOG.md index 90c47cb3..f5e3dec7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,22 @@ # Changelog -## Unreleased +## [0.12.1] - 2026-08-19 + +Published from the [wolkwork/obstore](https://github.com/wolkwork/obstore) fork as +`obstore-databaas` on PyPI while the `RemoteSignedS3Store` contribution is under review +upstream. The import name is unchanged, so this is a drop-in replacement for `obstore`. + +### New Features :magic_wand: + +- Add `RemoteSignedS3Store`, an S3-compatible store that never receives S3 credentials + and instead has every S3 REST request signed by a Python callback immediately before + the request is dispatched. This suits catalogs that vend signatures rather than + credentials, such as Lakekeeper's S3 request signer. Ranged `GET`, `HEAD`, `PUT` + (including conditional create/update, attributes and tags), `DELETE`, server-side + `COPY`, `LIST` and multipart uploads are supported, and every retry is signed afresh. + `RemoteSignedS3Store.from_s3_url(location, signer, endpoint=...)` builds a store from an + `s3://bucket/prefix` location plus the endpoint serving it, which is how catalogs + usually report a location. ## [0.11.0] - 2026-06-25 diff --git a/Cargo.lock b/Cargo.lock index a4d0d201..ff4bf947 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1464,7 +1464,7 @@ dependencies = [ [[package]] name = "obstore" -version = "0.11.0" +version = "0.12.1" dependencies = [ "arrow", "bytes", @@ -1736,6 +1736,8 @@ dependencies = [ "percent-encoding", "pyo3", "pyo3-async-runtimes", + "quick-xml", + "reqwest 0.13.4", "serde", "thiserror 1.0.69", "tokio", diff --git a/docs/api/store/remote-signed-s3.md b/docs/api/store/remote-signed-s3.md new file mode 100644 index 00000000..b46fc221 --- /dev/null +++ b/docs/api/store/remote-signed-s3.md @@ -0,0 +1,5 @@ +# Remote-signed S3 + +::: obstore.store.RemoteSignedS3Store + options: + inherited_members: true diff --git a/mkdocs.yml b/mkdocs.yml index 036c6c31..c77b2a11 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -51,6 +51,7 @@ nav: - obstore.store: - api/store/index.md - api/store/aws.md + - api/store/remote-signed-s3.md - api/store/gcs.md - api/store/azure.md - api/store/http.md diff --git a/obstore/Cargo.toml b/obstore/Cargo.toml index 06b59cda..3c659ea0 100644 --- a/obstore/Cargo.toml +++ b/obstore/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "obstore" -version = "0.11.0" +version = "0.12.1" authors = { workspace = true } edition = { workspace = true } description = "The simplest, highest-throughput interface to Amazon S3, Google Cloud Storage, Azure Blob Storage, and S3-compliant APIs like Cloudflare R2." diff --git a/obstore/pyproject.toml b/obstore/pyproject.toml index 714931ba..655e9b66 100644 --- a/obstore/pyproject.toml +++ b/obstore/pyproject.toml @@ -3,7 +3,7 @@ requires = ["maturin>=1.12.0,<2.0"] build-backend = "maturin" [project] -name = "obstore" +name = "obstore-databaas" requires-python = ">=3.10" readme = "README.md" dependencies = ["typing-extensions; python_version < '3.13'"] diff --git a/obstore/python/obstore/_store/__init__.pyi b/obstore/python/obstore/_store/__init__.pyi index 39eb60ec..0f2ab25a 100644 --- a/obstore/python/obstore/_store/__init__.pyi +++ b/obstore/python/obstore/_store/__init__.pyi @@ -21,6 +21,9 @@ from ._gcs import GCSCredential as GCSCredential from ._gcs import GCSCredentialProvider as GCSCredentialProvider from ._gcs import GCSStore as GCSStore from ._http import HTTPStore as HTTPStore +from ._remote_signed_s3 import RemoteSignedS3Store as RemoteSignedS3Store +from ._remote_signed_s3 import Signer as Signer +from ._remote_signed_s3 import SignerResult as SignerResult from ._retry import BackoffConfig as BackoffConfig from ._retry import RetryConfig as RetryConfig @@ -201,7 +204,13 @@ class MemoryStore: def __init__(self) -> None: ... ObjectStore: TypeAlias = ( - AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore + AzureStore + | GCSStore + | HTTPStore + | RemoteSignedS3Store + | S3Store + | LocalStore + | MemoryStore ) """All supported ObjectStore implementations. diff --git a/obstore/python/obstore/_store/_remote_signed_s3.pyi b/obstore/python/obstore/_store/_remote_signed_s3.pyi new file mode 100644 index 00000000..0245dbe7 --- /dev/null +++ b/obstore/python/obstore/_store/_remote_signed_s3.pyi @@ -0,0 +1,122 @@ +import sys +from collections.abc import Awaitable, Callable + +from ._client import ClientConfig +from ._retry import RetryConfig + +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +SignerResult: TypeAlias = tuple[str, dict[str, str]] +"""The signed `(uri, headers)` returned by a [`Signer`][obstore.store.Signer].""" + +Signer: TypeAlias = Callable[ + [str, str, dict[str, str]], + SignerResult | Awaitable[SignerResult], +] +"""A callback that signs a single S3 request. + +It is called with `(method, uri, headers)` immediately before each request is +dispatched, and must return the signed `(uri, headers)`. It may be synchronous or +asynchronous. + +The headers it returns **replace** the headers it was given, rather than being merged +into them, because a signature only covers the headers the signer chose to sign. A +signer must therefore echo back every header it was passed that the request still needs +(such as `range` or `content-length`) alongside the ones it adds. +""" + +class RemoteSignedS3Store: + """An S3-compatible store that has each request signed by a remote service.""" + + def __init__( + self, + url: str, + signer: Signer, + *, + virtual_hosted_style_request: bool = False, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + ) -> None: + """Construct a new RemoteSignedS3Store. + + Args: + url: Base URL including the bucket and optional object prefix. + signer: Callback that signs each request. + + Keyword Args: + virtual_hosted_style_request: Whether the bucket is named by the URL's host + rather than its first path segment. Defaults to `False`. + client_options: HTTP client options, such as timeouts and `allow_http`. + retry_config: How to retry failed requests. Every retry is signed again. + + """ + @classmethod + def from_s3_url( + cls, + url: str, + signer: Signer, + *, + endpoint: str, + virtual_hosted_style_request: bool = False, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + ) -> Self: + """Construct a store from an `s3://` location and the endpoint that serves it. + + Catalogs hand out locations as `s3://bucket/key` with the S3 endpoint configured + separately, so this saves you assembling the HTTPS URL yourself. + + ```py + store = RemoteSignedS3Store.from_s3_url( + "s3://warehouse/zarr/my-array", + signer, + endpoint="https://s3.eu-west-1.amazonaws.com", + ) + ``` + + Args: + url: An `s3://bucket/prefix` or `s3a://bucket/prefix` location. + signer: Callback that signs each request. + + Keyword Args: + endpoint: The `http://` or `https://` origin of the S3 endpoint serving the + bucket, such as `https://s3.eu-west-1.amazonaws.com`. It must not include + a path, since that would make the split between endpoint, bucket and key + prefix ambiguous; build the URL yourself in that case. + virtual_hosted_style_request: Set to `True` to address the bucket as a + subdomain of `endpoint`'s host rather than as its first path segment. + Defaults to `False`. + client_options: HTTP client options, such as timeouts and `allow_http`. + retry_config: How to retry failed requests. Every retry is signed again. + + """ + def __eq__(self, other: object) -> bool: ... + @property + def url(self) -> str: + """The base URL this store was constructed with.""" + @property + def bucket(self) -> str: + """The bucket name, taken from `url`'s first path segment or its host.""" + @property + def prefix(self) -> str | None: + """The key prefix implied by `url`, or `None` if it names only the bucket.""" + @property + def signer(self) -> Signer: + """The signer callback passed to the constructor.""" + @property + def virtual_hosted_style_request(self) -> bool: + """Whether the bucket is named by the URL's host.""" + @property + def client_options(self) -> ClientConfig | None: + """Get the store's client configuration.""" + @property + def retry_config(self) -> RetryConfig | None: + """Get the store's retry configuration.""" diff --git a/obstore/python/obstore/store.py b/obstore/python/obstore/store.py index 72fe1b80..3259016b 100644 --- a/obstore/python/obstore/store.py +++ b/obstore/python/obstore/store.py @@ -89,6 +89,7 @@ "HTTPStore", "LocalStore", "MemoryStore", + "RemoteSignedS3Store", "RetryConfig", "S3Config", "S3Credential", @@ -655,6 +656,125 @@ class MemoryStore(ObjectStoreMethods, _store.MemoryStore): """ +class RemoteSignedS3Store(ObjectStoreMethods, _store.RemoteSignedS3Store): + """An S3-compatible store that has each request signed by a remote service. + + Unlike [`S3Store`][obstore.store.S3Store], this store never receives or holds S3 + credentials. Instead, it constructs each final S3 REST request, hands the request's + method, URI, and headers to a `signer` callback immediately before dispatch, and + sends exactly the URI and headers the callback returns. The store re-signs on every + retry, so signatures may be short-lived. + + This suits setups where an external service (for example [Lakekeeper]'s S3 request + signer) authorizes access to a location without vending credentials to the client. + + The `signer` contract is: + + ```py + def signer( + method: str, + uri: str, + headers: dict[str, str], + ) -> tuple[str, dict[str, str]]: + ... + ``` + + It receives the HTTP method, the fully-constructed request URI (including any query + parameters), and the request headers, and returns the signed `(uri, headers)`. The + callback may be synchronous or asynchronous; any credential/token caching belongs + inside the callback. + + !!! note + A synchronous `signer` is required when calling the synchronous store methods + (for example `get_range` or `list().collect()`). Use an asynchronous `signer` + with the `_async` methods. + + !!! important "The signer's headers replace, not extend, the request's" + The store sends exactly the headers the callback returns. This matches the S3 + remote-signing contract: a signature covers a specific set of headers, so + keeping headers the signer did not return would invalidate it. A signer must + therefore echo back every header it was passed that the request still needs — + `range`, `content-length`, `if-none-match` and so on — alongside the ones it + adds. + + The request body is never sent to the signer, so a signer cannot compute a + payload hash and must sign with `x-amz-content-sha256: UNSIGNED-PAYLOAD`. + + `GET` (including byte ranges), `HEAD`, `PUT` (with conditional create/overwrite and + object attributes/tags), `DELETE`, server-side `COPY`, `LIST` and multipart uploads + are all supported. Every multipart initiation, part upload, completion and abort is + an independently signed request. + + !!! note "Plain HTTP endpoints" + Like the other stores, this one refuses `http://` URLs unless you opt in with + `client_options={"allow_http": True}`. + + !!! warning "`LIST` against a path-validating signer" + `LIST` uses S3 `ListObjectsV2`, which issues a request to the bucket root + with the prefix as a *query parameter* + (`GET /bucket?list-type=2&prefix=...`). Signers that authorize by the URL + *path* rather than the query — notably Lakekeeper's generic-table signer — + cannot match the location and reject these requests (e.g. + `NoSuchTableLocationException`). + + This affects any operation that lists, including Zarr's `delete_dir` + (triggered by `overwrite=True`). With such a signer, avoid list-based + operations against the signed location: write arrays fresh instead of + overwriting. There is no client-side fix — the prefix cannot be moved into + the URL path without breaking S3 list semantics. + + **Example**: + + ```py + import requests + from obstore.store import RemoteSignedS3Store + + def signer(method, uri, headers): + resp = requests.post( + "https://catalog.example.com/sign", + json={"method": method, "uri": uri, "headers": headers}, + ) + resp.raise_for_status() + signed = resp.json() + return signed["uri"], signed["headers"] + + store = RemoteSignedS3Store("https://s3.example.com/bucket/prefix", signer) + data = store.get_range("chunk.bin", start=0, end=1024) + ``` + + Catalogs usually hand out an `s3://` location with the endpoint configured + separately, in which case + [`from_s3_url`][obstore.store.RemoteSignedS3Store.from_s3_url] saves assembling the + URL: + + ```py + store = RemoteSignedS3Store.from_s3_url( + "s3://warehouse/zarr/my-array", + signer, + endpoint="https://s3.example.com", + ) + ``` + + [Lakekeeper]: https://lakekeeper.io/ + + Args: + url: Base URL including the bucket and optional object prefix, for example + `https://s3.example.com/bucket/prefix`. All paths passed to store methods + are resolved relative to this prefix. + signer: Callable that signs each request, as described above. + + Keyword Args: + virtual_hosted_style_request: Set to `True` when the bucket is named by the + URL's host rather than its first path segment, as in + `https://bucket.s3.amazonaws.com/prefix`. Every path segment is then treated + as part of the key prefix. Defaults to `False`. + client_options: HTTP client options, such as timeouts and `allow_http`. + retry_config: How to retry failed requests. Each attempt is signed afresh, so a + retry never reuses an expired signature. + + """ + + class S3Store(ObjectStoreMethods, _store.S3Store): """Interface to an Amazon S3 bucket. @@ -679,6 +799,7 @@ class S3Store(ObjectStoreMethods, _store.S3Store): AzureStore, GCSStore, HTTPStore, + RemoteSignedS3Store, S3Store, LocalStore, MemoryStore, diff --git a/pyo3-object_store/CHANGELOG.md b/pyo3-object_store/CHANGELOG.md index 127705ba..e2631721 100644 --- a/pyo3-object_store/CHANGELOG.md +++ b/pyo3-object_store/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## Unreleased + +- Add `PyRemoteSignedS3Store` (in `aws::remote_signed`), exporting a + `RemoteSignedS3Store` that signs every S3 request through a Python callback instead of + holding S3 credentials. Like `PyS3Store`, it applies its key prefix with + `MaybePrefixedStore`. + ## [0.12.0] - 2026-06-25 - Bump to object_store 0.14 diff --git a/pyo3-object_store/Cargo.toml b/pyo3-object_store/Cargo.toml index 9c6577da..452425a1 100644 --- a/pyo3-object_store/Cargo.toml +++ b/pyo3-object_store/Cargo.toml @@ -22,6 +22,8 @@ async-trait = "0.1.85" bytes = "1" chrono = "0.4" futures = "0.3" +reqwest = { version = "0.13", default-features = false, features = ["rustls"] } +quick-xml = { version = "0.41", features = ["serialize"] } # This is already an object_store dependency humantime = "2.1" # This is already an object_store dependency @@ -38,9 +40,9 @@ object_store = { version = "0.14.0", features = [ percent-encoding = "2.1" pyo3 = { version = "0.29", features = ["chrono", "indexmap"] } pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } -serde = "1" +serde = { version = "1", features = ["derive"] } thiserror = "1" -tokio = { version = "1.40", features = ["rt-multi-thread"] } +tokio = { version = "1.40", features = ["rt-multi-thread", "time"] } url = "2" [lib] diff --git a/pyo3-object_store/src/api.rs b/pyo3-object_store/src/api.rs index ad22d32f..88bc1bd9 100644 --- a/pyo3-object_store/src/api.rs +++ b/pyo3-object_store/src/api.rs @@ -3,7 +3,8 @@ use pyo3::prelude::*; use crate::error::*; use crate::{ - from_url, PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store, + from_url, PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, + PyRemoteSignedS3Store, PyS3Store, }; /// Export the default Python API as a submodule named `store` within the given parent module @@ -55,6 +56,7 @@ pub fn register_store_module( child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; // Set the value of `__module__` correctly on each publicly exposed function or class let __module__ = intern!(py, "__module__"); @@ -79,6 +81,9 @@ pub fn register_store_module( child_module .getattr("S3Store")? .setattr(__module__, &full_module_string)?; + child_module + .getattr("RemoteSignedS3Store")? + .setattr(__module__, &full_module_string)?; // Add the child module to the parent module parent_module.add_submodule(&child_module)?; diff --git a/pyo3-object_store/src/aws/mod.rs b/pyo3-object_store/src/aws/mod.rs index a4871db1..a34bc144 100644 --- a/pyo3-object_store/src/aws/mod.rs +++ b/pyo3-object_store/src/aws/mod.rs @@ -1,4 +1,6 @@ mod credentials; +mod remote_signed; mod store; +pub use remote_signed::PyRemoteSignedS3Store; pub use store::PyS3Store; diff --git a/pyo3-object_store/src/aws/remote_signed/client.rs b/pyo3-object_store/src/aws/remote_signed/client.rs new file mode 100644 index 00000000..1fd39136 --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/client.rs @@ -0,0 +1,229 @@ +//! Turning an operation into a signed S3 REST request, and dispatching it with retries. + +use std::time::{Duration, Instant}; + +use http::header::IF_NONE_MATCH; +use http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode}; +use object_store::client::{HttpError, HttpErrorKind, HttpRequestBody, HttpResponse}; +use object_store::path::Path; +use object_store::{PutPayload, Result}; +use percent_encoding::{utf8_percent_encode, AsciiSet, NON_ALPHANUMERIC}; +use serde::Deserialize; +use url::Url; + +use super::{error_for, RemoteSignedS3Store, STORE}; + +/// SigV4 requires the unreserved characters of RFC 3986 to be left alone and everything else to be +/// percent-encoded. `/` is excluded because it separates key segments. +/// +/// +const KEY_ENCODE_SET: AsciiSet = NON_ALPHANUMERIC + .remove(b'-') + .remove(b'.') + .remove(b'_') + .remove(b'~') + .remove(b'/'); + +/// One S3 REST request, before signing. +/// +/// Held rather than dispatched directly so that [`RemoteSignedS3Store::send`] can re-sign and +/// resend it on retry. Signatures are short-lived, so a retry must never reuse the previous +/// signature. +pub(super) struct SignedRequest { + method: Method, + url: Url, + headers: HeaderMap, + body: PutPayload, + /// Whether this request carries `If-None-Match: *`, in which case a rejected precondition + /// means "the object already exists" rather than a generic failure. + conditional_create: bool, +} + +impl SignedRequest { + pub(super) fn new(method: Method, url: Url) -> Self { + Self { + method, + url, + headers: HeaderMap::new(), + body: PutPayload::default(), + conditional_create: false, + } + } + + pub(super) fn header(mut self, name: impl Into, value: HeaderValue) -> Self { + self.headers.insert(name.into(), value); + self + } + + pub(super) fn headers(mut self, headers: HeaderMap) -> Self { + self.headers.extend(headers); + self + } + + pub(super) fn body(mut self, body: PutPayload) -> Self { + self.body = body; + self + } + + pub(super) fn conditional_create(mut self) -> Self { + self.headers + .insert(IF_NONE_MATCH, HeaderValue::from_static("*")); + self.conditional_create = true; + self + } +} + +impl RemoteSignedS3Store { + /// The URL of `location`, with the key percent-encoded as SigV4 requires. + /// + /// `Url::join` cannot be used here: it would interpret `?` and `#` in a key as the start of a + /// query or fragment, silently addressing a different object. + pub(super) fn object_url(&self, location: &Path) -> Result { + self.bucket_url + .join(&encode_key(location)) + .map_err(|error| error_for("invalid object URL", error)) + } + + pub(super) fn object_request(&self, method: Method, location: &Path) -> Result { + Ok(SignedRequest::new(method, self.object_url(location)?)) + } + + /// Sign and dispatch `request`, retrying transient failures. + /// + /// Each attempt is signed afresh, so a retry never depends on the lifetime of an earlier + /// signature. + pub(super) async fn send( + &self, + location: &Path, + request: SignedRequest, + ) -> Result { + let deadline = Instant::now() + self.retry.retry_timeout; + let mut backoff = self.retry.backoff.init_backoff; + let mut attempts = 0; + loop { + let (uri, headers) = self + .signer + .sign(&request.method, &request.url, &request.headers) + .await?; + let mut builder = http::Request::builder() + .method(request.method.clone()) + .uri(uri.as_str()); + if let Some(target) = builder.headers_mut() { + *target = headers; + } + let http_request = builder + .body(HttpRequestBody::from(request.body.clone())) + .map_err(|error| error_for("invalid HTTP request", error))?; + + let result = self.client.execute(http_request).await; + let retryable = match &result { + Err(error) => is_retryable_error(error), + Ok(response) => is_retryable_status(response.status()), + }; + if !retryable || attempts >= self.retry.max_retries || Instant::now() >= deadline { + let response = result.map_err(|error| error_for("HTTP request failed", error))?; + return check_response(location, request.conditional_create, response).await; + } + + tokio::time::sleep(backoff).await; + backoff = Duration::min( + backoff.mul_f64(self.retry.backoff.base), + self.retry.backoff.max_backoff, + ); + attempts += 1; + } + } +} + +/// Percent-encode a key for use in a request path, as SigV4 requires. +pub(super) fn encode_key(location: &Path) -> String { + utf8_percent_encode(location.as_ref(), &KEY_ENCODE_SET).to_string() +} + +/// The `` document S3 returns for a failed request. +#[derive(Deserialize)] +#[serde(rename = "Error", rename_all = "PascalCase")] +struct S3Error { + code: Option, + message: Option, +} + +/// Whether a transport-level failure is worth retrying. +fn is_retryable_error(error: &HttpError) -> bool { + matches!( + error.kind(), + HttpErrorKind::Connect + | HttpErrorKind::Request + | HttpErrorKind::Timeout + | HttpErrorKind::Interrupted + ) +} + +/// Whether a response status is worth retrying. +/// +/// Retrying cannot corrupt an object: writing the same key or part twice is idempotent, and this +/// matches what `object_store`'s own S3 client retries. +/// +/// Two cases are lossy rather than harmful, and are accepted for the same reason `object_store` +/// accepts them. If a `CompleteMultipartUpload` succeeds but its response is lost, the retry sees +/// `NoSuchUpload` and reports `NotFound` even though the object exists. If a conditional create +/// (`If-None-Match: *`) succeeds but its response is lost, the retry sees a rejected precondition +/// and reports `AlreadyExists`, as though another writer had won the race. +fn is_retryable_status(status: StatusCode) -> bool { + status.is_server_error() + || matches!( + status, + StatusCode::REQUEST_TIMEOUT | StatusCode::TOO_MANY_REQUESTS + ) +} + +/// Turn a non-2xx response into the matching `object_store` error, including whatever detail S3 +/// put in the response body. +async fn check_response( + location: &Path, + conditional_create: bool, + response: HttpResponse, +) -> Result { + let status = response.status(); + if status.is_success() { + return Ok(response); + } + let body = response.into_body().bytes().await.unwrap_or_default(); + let source: Box = match s3_error_detail(&body) { + Some(detail) => format!("HTTP {status}: {detail}").into(), + None => format!("HTTP {status}").into(), + }; + let path = location.to_string(); + // A conditional create that lost the race reports a rejected precondition, but callers such + // as Zarr's `set_if_not_exists` expect `AlreadyExists`. S3 uses 412; some implementations + // report the same race as 409. + if conditional_create + && matches!( + status, + StatusCode::PRECONDITION_FAILED | StatusCode::CONFLICT + ) + { + return Err(object_store::Error::AlreadyExists { path, source }); + } + Err(match status { + StatusCode::NOT_FOUND => object_store::Error::NotFound { path, source }, + StatusCode::UNAUTHORIZED => object_store::Error::Unauthenticated { path, source }, + StatusCode::FORBIDDEN => object_store::Error::PermissionDenied { path, source }, + StatusCode::PRECONDITION_FAILED => object_store::Error::Precondition { path, source }, + StatusCode::NOT_MODIFIED => object_store::Error::NotModified { path, source }, + _ => object_store::Error::Generic { + store: STORE, + source, + }, + }) +} + +/// The `Code: Message` of an S3 `` body, if `body` is one. +pub(super) fn s3_error_detail(body: &[u8]) -> Option { + let error: S3Error = quick_xml::de::from_reader(body).ok()?; + match (error.code, error.message) { + (Some(code), Some(message)) => Some(format!("{code}: {message}")), + (Some(detail), None) | (None, Some(detail)) => Some(detail), + (None, None) => None, + } +} diff --git a/pyo3-object_store/src/aws/remote_signed/headers.rs b/pyo3-object_store/src/aws/remote_signed/headers.rs new file mode 100644 index 00000000..9df667ba --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/headers.rs @@ -0,0 +1,174 @@ +//! Translating between S3 headers and the `object_store` types that carry them. + +use std::ops::Range; + +use chrono::{DateTime, Utc}; +use http::header::{ + CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_LENGTH, + CONTENT_RANGE, CONTENT_TYPE, ETAG, LAST_MODIFIED, +}; +use http::{HeaderMap, HeaderName, HeaderValue, StatusCode}; +use object_store::path::Path; +use object_store::{Attribute, AttributeValue, Attributes, ObjectMeta, Result, TagSet}; + +use super::{error_for, STORE}; + +/// The `x-amz-*` headers this store sets or reads itself. +pub(super) const STORAGE_CLASS: &str = "x-amz-storage-class"; +pub(super) const TAGGING: &str = "x-amz-tagging"; +pub(super) const USER_METADATA_PREFIX: &str = "x-amz-meta-"; +pub(super) const VERSION_ID: &str = "x-amz-version-id"; + +/// The HTTP date format used by the `If-[Un]Modified-Since` headers. +pub(super) const HTTP_DATE_FORMAT: &str = "%a, %d %b %Y %H:%M:%S GMT"; + +/// Map the `Attributes` obstore accepts onto the S3 headers that carry them. +pub(super) fn attribute_headers(attributes: &Attributes) -> Result { + let mut headers = HeaderMap::new(); + for (attribute, value) in attributes.iter() { + let name = match attribute { + Attribute::CacheControl => CACHE_CONTROL, + Attribute::ContentDisposition => CONTENT_DISPOSITION, + Attribute::ContentEncoding => CONTENT_ENCODING, + Attribute::ContentLanguage => CONTENT_LANGUAGE, + Attribute::ContentType => CONTENT_TYPE, + Attribute::StorageClass => HeaderName::from_static(STORAGE_CLASS), + Attribute::Metadata(key) => { + HeaderName::from_bytes(format!("{USER_METADATA_PREFIX}{key}").as_bytes()) + .map_err(|error| error_for("invalid metadata key", error))? + } + // `Attribute` is `#[non_exhaustive]`, so a future variant must be rejected rather + // than silently dropped: callers set attributes expecting them to be stored. + other => { + return Err(object_store::Error::NotSupported { + source: format!("attribute {other:?} is not supported by {STORE}").into(), + }) + } + }; + headers.insert( + name, + HeaderValue::from_str(value.as_ref()) + .map_err(|error| error_for("invalid attribute value", error))?, + ); + } + Ok(headers) +} + +pub(super) fn tag_headers(tags: &TagSet) -> Result { + let mut headers = HeaderMap::new(); + if !tags.is_empty() { + headers.insert( + HeaderName::from_static(TAGGING), + HeaderValue::from_str(tags.encoded()) + .map_err(|error| error_for("invalid tag set", error))?, + ); + } + Ok(headers) +} + +/// Recover the attributes of an object from its response headers. +pub(super) fn response_attributes(headers: &HeaderMap) -> Attributes { + let mut attributes = Attributes::new(); + for (attribute, name) in [ + (Attribute::CacheControl, CACHE_CONTROL), + (Attribute::ContentDisposition, CONTENT_DISPOSITION), + (Attribute::ContentEncoding, CONTENT_ENCODING), + (Attribute::ContentLanguage, CONTENT_LANGUAGE), + (Attribute::ContentType, CONTENT_TYPE), + ( + Attribute::StorageClass, + HeaderName::from_static(STORAGE_CLASS), + ), + ] { + if let Some(value) = header_str(headers, &name) { + attributes.insert(attribute, AttributeValue::from(value)); + } + } + for (name, value) in headers.iter() { + if let Some(key) = name.as_str().strip_prefix(USER_METADATA_PREFIX) { + if let Ok(value) = value.to_str() { + attributes.insert( + Attribute::Metadata(key.to_owned().into()), + AttributeValue::from(value.to_owned()), + ); + } + } + } + attributes +} + +pub(super) fn header_str(headers: &HeaderMap, name: &HeaderName) -> Option { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned) +} + +pub(super) fn object_meta( + location: &Path, + headers: &HeaderMap, + fallback_size: u64, +) -> Result { + let size = headers + .get(CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .map(str::parse) + .transpose() + .map_err(|error| error_for("invalid Content-Length", error))? + .unwrap_or(fallback_size); + let last_modified = headers + .get(LAST_MODIFIED) + .and_then(|value| value.to_str().ok()) + .map(DateTime::parse_from_rfc2822) + .transpose() + .map_err(|error| error_for("invalid Last-Modified", error))? + .map(|value| value.with_timezone(&Utc)) + .unwrap_or(DateTime::::UNIX_EPOCH); + Ok(ObjectMeta { + location: location.clone(), + last_modified, + size, + e_tag: header_str(headers, &ETAG), + version: header_str(headers, &HeaderName::from_static(VERSION_ID)), + }) +} + +/// The `(range, total size)` of a `Content-Range: bytes -/` header. +pub(super) fn content_range(headers: &HeaderMap) -> Option<(Range, u64)> { + let value = headers + .get(CONTENT_RANGE)? + .to_str() + .ok()? + .strip_prefix("bytes ")?; + let (range, size) = value.split_once('/')?; + let (start, end) = range.split_once('-')?; + Some(( + start.parse().ok()?..end.parse::().ok()?.checked_add(1)?, + size.parse().ok()?, + )) +} + +pub(super) fn response_range( + headers: &HeaderMap, + status: StatusCode, + meta: &ObjectMeta, + requested: Option<&object_store::GetRange>, +) -> Result> { + if status == StatusCode::PARTIAL_CONTENT { + return content_range(headers) + .map(|(range, _)| range) + .ok_or_else(|| { + error_for( + "invalid Content-Range", + std::io::Error::other("missing or malformed header"), + ) + }); + } + if requested.is_some() { + return Err(error_for( + "range request did not return 206", + std::io::Error::other("unexpected response status"), + )); + } + Ok(0..meta.size) +} diff --git a/pyo3-object_store/src/aws/remote_signed/list.rs b/pyo3-object_store/src/aws/remote_signed/list.rs new file mode 100644 index 00000000..e7d8d64c --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/list.rs @@ -0,0 +1,180 @@ +//! `ListObjectsV2` requests, their XML response, and paging through them. + +use std::borrow::Cow; + +use chrono::{DateTime, Utc}; +use futures::stream::{self, BoxStream}; +use futures::{StreamExt, TryStreamExt}; +use http::Method; +use object_store::path::Path; +use object_store::{ObjectMeta, Result}; +use percent_encoding::percent_decode_str; +use serde::Deserialize; + +use super::client::SignedRequest; +use super::{error_for, RemoteSignedS3Store}; + +impl RemoteSignedS3Store { + pub(super) async fn list_page( + &self, + prefix: Option<&Path>, + delimiter: Option<&str>, + offset: Option<&Path>, + token: Option<&str>, + ) -> Result { + let mut url = self.bucket_url.clone(); + { + let mut query = url.query_pairs_mut(); + query.append_pair("list-type", "2"); + if let Some(prefix) = format_prefix(prefix) { + query.append_pair("prefix", &prefix); + } + // Ask S3 to percent-encode keys so that control characters, which are legal in a key + // but not in XML, survive the response. + query.append_pair("encoding-type", "url"); + if let Some(delimiter) = delimiter { + query.append_pair("delimiter", delimiter); + } + if let Some(offset) = offset { + query.append_pair("start-after", offset.as_ref()); + } + if let Some(token) = token { + query.append_pair("continuation-token", token); + } + } + let response = self + .send(&Path::default(), SignedRequest::new(Method::GET, url)) + .await?; + let body = response + .into_body() + .bytes() + .await + .map_err(|error| error_for("failed to read list response", error))?; + let response: ListObjectsV2 = quick_xml::de::from_reader(body.as_ref()) + .map_err(|error| error_for("invalid ListObjectsV2 response", error))?; + let objects = response + .contents + .into_iter() + .map(ListObject::into_meta) + .collect::>>()?; + let common_prefixes = response + .common_prefixes + .into_iter() + .map(|prefix| Path::parse(decode_key(&prefix.prefix))) + .collect::, _>>()?; + Ok(ListPage { + objects, + common_prefixes, + next_token: response.next_continuation_token, + }) + } + + /// Stream every object under `prefix`, fetching each `ListObjectsV2` page only once the + /// previous page has been consumed. + pub(super) fn list_paginated( + &self, + prefix: Option<&Path>, + offset: Option<&Path>, + ) -> BoxStream<'static, Result> { + let store = self.clone(); + let prefix = prefix.cloned(); + let offset = offset.cloned(); + stream::try_unfold(Some(None::), move |token| { + let store = store.clone(); + let prefix = prefix.clone(); + let offset = offset.clone(); + async move { + // `None` marks the stream as exhausted; `Some(None)` is the first request. + let Some(token) = token else { + return Ok::<_, object_store::Error>(None); + }; + let page = store + .list_page(prefix.as_ref(), None, offset.as_ref(), token.as_deref()) + .await?; + let next = page.next_token.map(Some); + Ok(Some((stream::iter(page.objects.into_iter().map(Ok)), next))) + } + }) + .try_flatten() + .boxed() + } +} + +#[derive(Deserialize)] +#[serde(rename = "ListBucketResult")] +struct ListObjectsV2 { + #[serde(rename = "Contents", default)] + contents: Vec, + #[serde(rename = "CommonPrefixes", default)] + common_prefixes: Vec, + #[serde(rename = "NextContinuationToken")] + next_continuation_token: Option, +} + +#[derive(Deserialize)] +struct ListObject { + #[serde(rename = "Key")] + key: String, + #[serde(rename = "LastModified")] + last_modified: String, + #[serde(rename = "Size")] + size: u64, + #[serde(rename = "ETag")] + e_tag: Option, +} + +impl ListObject { + fn into_meta(self) -> Result { + Ok(ObjectMeta { + location: Path::parse(decode_key(&self.key))?, + last_modified: DateTime::parse_from_rfc3339(&self.last_modified) + .map_err(|error| error_for("invalid list LastModified", error))? + .with_timezone(&Utc), + size: self.size, + e_tag: self.e_tag, + version: None, + }) + } +} + +#[derive(Deserialize)] +struct ListPrefix { + #[serde(rename = "Prefix")] + prefix: String, +} + +pub(super) struct ListPage { + pub(super) objects: Vec, + pub(super) common_prefixes: Vec, + pub(super) next_token: Option, +} + +/// The `prefix` query parameter for a `ListObjectsV2` request. +/// +/// A trailing delimiter is appended so that the prefix matches on segment boundaries: without +/// it, listing `a/b` would also return `a/bc/d`. This mirrors what `object_store`'s own list +/// clients do. +fn format_prefix(prefix: Option<&Path>) -> Option { + prefix + .filter(|prefix| !prefix.as_ref().is_empty()) + .map(|prefix| format!("{}/", prefix.as_ref())) +} + +/// Build a key prefix out of the path segments of a store URL. +/// +/// A `Url`'s segments are percent-encoded, but an `object_store::Path` holds a *decoded* key and +/// [`encode_key`] re-encodes it on the way out. Without decoding here, a URL ending in `caf%C3%A9` +/// would produce the literal seven-character prefix `caf%C3%A9`, be encoded again as +/// `caf%25C3%25A9`, and address the wrong object. +pub(super) fn decode_prefix<'a>(segments: impl Iterator) -> String { + segments.map(decode_key).collect::>().join("/") +} + +/// Undo the `encoding-type=url` encoding of a listed key, leaving it unchanged if it is not valid +/// percent-encoded UTF-8. +fn decode_key(key: &str) -> String { + percent_decode_str(key) + .decode_utf8() + .map(Cow::into_owned) + .unwrap_or_else(|_| key.to_owned()) +} diff --git a/pyo3-object_store/src/aws/remote_signed/mod.rs b/pyo3-object_store/src/aws/remote_signed/mod.rs new file mode 100644 index 00000000..37dc2428 --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/mod.rs @@ -0,0 +1,292 @@ +//! An S3-compatible store that has every request signed by a Python callback. +//! +//! The work is split so each concern stays readable on its own: +//! +//! - [`signer`] is the Python boundary, +//! - [`client`] builds, signs, dispatches and retries one request, +//! - [`list`], [`multipart`] and [`headers`] cover the S3 protocol details, +//! - [`store`] is the Python-facing class, +//! - and this module holds the store type and its [`ObjectStore`] implementation. + +mod client; +mod headers; +mod list; +mod multipart; +mod signer; +mod store; + +use std::fmt::{Display, Formatter}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use futures::stream::{self, BoxStream}; +use futures::{StreamExt, TryStreamExt}; +use http::header::{ + CONTENT_LENGTH, ETAG, IF_MATCH, IF_MODIFIED_SINCE, IF_NONE_MATCH, IF_UNMODIFIED_SINCE, RANGE, +}; +use http::{HeaderName, HeaderValue, Method}; +use object_store::client::HttpClient; +use object_store::path::Path; +use object_store::{ + CopyMode, CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult, + Result, RetryConfig, +}; +use url::Url; + +pub use store::PyRemoteSignedS3Store; + +use client::{encode_key, SignedRequest}; +use headers::{ + attribute_headers, content_range, header_str, object_meta, response_attributes, response_range, + tag_headers, HTTP_DATE_FORMAT, VERSION_ID, +}; +use multipart::{MultipartState, RemoteSignedMultipartUpload}; +use signer::PySigner; + +const STORE: &str = "RemoteSignedS3"; + +/// The `x-amz-copy-source` header, which always names the bucket. +const COPY_SOURCE: &str = "x-amz-copy-source"; + +/// An S3-compatible object store that obtains a signature for every HTTP request from Python. +/// +/// This addresses the bucket root and takes whole keys. Any key prefix is applied by wrapping +/// it in a [`MaybePrefixedStore`], exactly as [`PyS3Store`](super::PyS3Store) does. +#[derive(Debug, Clone)] +pub struct RemoteSignedS3Store { + /// The bucket root that keys are resolved against. Always ends in `/`. + bucket_url: Url, + /// The bucket name. Needed by `x-amz-copy-source`, which always names the bucket even + /// when the request URL does not. + bucket: String, + signer: PySigner, + client: HttpClient, + retry: RetryConfig, +} + +impl Display for RemoteSignedS3Store { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{STORE}({})", self.bucket_url) + } +} + +#[async_trait] +impl ObjectStore for RemoteSignedS3Store { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + let mut request = self + .object_request(Method::PUT, location)? + .header( + CONTENT_LENGTH, + HeaderValue::from_str(&payload.content_length().to_string()) + .expect("content length is a valid header value"), + ) + .headers(attribute_headers(&opts.attributes)?) + .headers(tag_headers(&opts.tags)?) + .body(payload); + match opts.mode { + PutMode::Overwrite => {} + PutMode::Create => request = request.conditional_create(), + PutMode::Update(version) => { + let e_tag = version + .e_tag + .ok_or_else(|| object_store::Error::NotSupported { + source: "RemoteSignedS3Store requires an ETag to update a specific version" + .into(), + })?; + request = request.header( + IF_MATCH, + HeaderValue::from_str(&e_tag) + .map_err(|error| error_for("invalid ETag", error))?, + ); + } + } + let response = self.send(location, request).await?; + Ok(PutResult { + e_tag: header_str(response.headers(), &ETAG), + version: header_str(response.headers(), &HeaderName::from_static(VERSION_ID)), + extensions: Default::default(), + }) + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + let upload_id = self.create_multipart(location, opts).await?; + Ok(Box::new(RemoteSignedMultipartUpload { + state: Arc::new(MultipartState { + store: Arc::new(self.clone()), + location: location.clone(), + upload_id, + parts: Mutex::new(Vec::new()), + }), + part_idx: 0, + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + let mut url = self.object_url(location)?; + if let Some(version) = &options.version { + url.query_pairs_mut().append_pair("versionId", version); + } + let mut request = SignedRequest::new( + if options.head { + Method::HEAD + } else { + Method::GET + }, + url, + ); + if let Some(range) = &options.range { + request = request.header( + RANGE, + HeaderValue::from_str(&range.to_string()).expect("range is a valid header value"), + ); + } + for (name, value) in [ + (IF_MATCH, options.if_match.clone()), + (IF_NONE_MATCH, options.if_none_match.clone()), + ( + IF_MODIFIED_SINCE, + options + .if_modified_since + .map(|date| date.format(HTTP_DATE_FORMAT).to_string()), + ), + ( + IF_UNMODIFIED_SINCE, + options + .if_unmodified_since + .map(|date| date.format(HTTP_DATE_FORMAT).to_string()), + ), + ] { + if let Some(value) = value { + request = request.header( + name.clone(), + HeaderValue::from_str(&value) + .map_err(|error| error_for(name.as_str(), error))?, + ); + } + } + + let response = self.send(location, request).await?; + let status = response.status(); + let headers = response.headers().clone(); + + let mut meta = object_meta(location, &headers, 0)?; + if let Some((_, size)) = content_range(&headers) { + meta.size = size; + } + options.check_preconditions(&meta)?; + let range = response_range(&headers, status, &meta, options.range.as_ref())?; + let attributes = response_attributes(&headers); + + // Stream the body rather than collecting it, so that reading a large object does not + // require buffering the whole object in memory. + let payload = if options.head { + GetResultPayload::Stream(stream::empty().boxed()) + } else { + GetResultPayload::Stream( + response + .into_body() + .bytes_stream() + .map_err(|error| error_for("failed to read response body", error)) + .boxed(), + ) + }; + Ok(GetResult { + payload, + meta, + range, + attributes, + extensions: Default::default(), + }) + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + let store = Arc::new(self.clone()); + locations + .map(move |location| { + let store = Arc::clone(&store); + async move { + let location = location?; + let request = store.object_request(Method::DELETE, &location)?; + store.send(&location, request).await?; + Ok(location) + } + }) + // S3 has no signable bulk-delete (its `POST ?delete` needs a body checksum the signer + // never sees), so delete concurrently instead, as the S3 store does. + .buffered(10) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.list_paginated(prefix, None) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + self.list_paginated(prefix, Some(offset)) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + let mut objects = Vec::new(); + let mut common_prefixes = Vec::new(); + let mut token = None; + loop { + let page = self + .list_page(prefix, Some("/"), None, token.as_deref()) + .await?; + objects.extend(page.objects); + common_prefixes.extend(page.common_prefixes); + match page.next_token { + Some(next) => token = Some(next), + None => break, + } + } + Ok(ListResult { + common_prefixes, + objects, + extensions: Default::default(), + }) + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { + // `x-amz-copy-source` is `//` regardless of addressing style, so it is + // built from the bucket name rather than from the request URL's path. + let source = format!("/{}/{}", self.bucket, encode_key(from)); + let mut request = self.object_request(Method::PUT, to)?.header( + HeaderName::from_static(COPY_SOURCE), + HeaderValue::from_str(&source) + .map_err(|error| error_for("invalid copy source", error))?, + ); + if options.mode == CopyMode::Create { + request = request.conditional_create(); + } + self.send(to, request).await?; + Ok(()) + } +} + +fn error_for( + message: &str, + source: impl std::error::Error + Send + Sync + 'static, +) -> object_store::Error { + object_store::Error::Generic { + store: STORE, + source: format!("{message}: {source}").into(), + } +} diff --git a/pyo3-object_store/src/aws/remote_signed/multipart.rs b/pyo3-object_store/src/aws/remote_signed/multipart.rs new file mode 100644 index 00000000..8684df35 --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/multipart.rs @@ -0,0 +1,232 @@ +//! Multipart upload: initiation, parts, completion and abort, each signed independently. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use futures::FutureExt; +use http::header::{CONTENT_LENGTH, ETAG}; +use http::{HeaderName, HeaderValue, Method}; +use object_store::multipart::PartId; +use object_store::path::Path; +use object_store::{ + MultipartUpload, PutMultipartOptions, PutPayload, PutResult, Result, UploadPart, +}; +use serde::{Deserialize, Serialize}; + +use super::client::{s3_error_detail, SignedRequest}; +use super::headers::{header_str, VERSION_ID}; +use super::{attribute_headers, error_for, tag_headers, RemoteSignedS3Store, STORE}; + +impl RemoteSignedS3Store { + pub(super) async fn create_multipart( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result { + let PutMultipartOptions { + tags, + attributes, + extensions: _, + } = opts; + let mut url = self.object_url(location)?; + url.set_query(Some("uploads=")); + let request = SignedRequest::new(Method::POST, url) + .header(CONTENT_LENGTH, HeaderValue::from_static("0")) + .headers(attribute_headers(&attributes)?) + .headers(tag_headers(&tags)?); + let body = self + .send(location, request) + .await? + .into_body() + .bytes() + .await + .map_err(|error| error_for("failed to read CreateMultipartUpload response", error))?; + let response: InitiateMultipartUploadResult = quick_xml::de::from_reader(body.as_ref()) + .map_err(|error| error_for("invalid CreateMultipartUpload response", error))?; + Ok(response.upload_id) + } +} + +/// The shared state of an in-flight multipart upload. +/// +/// `MultipartUpload::put_part` must return a `'static` future so that parts can be uploaded +/// concurrently, which is why the store, upload id and completed parts are all held behind an +/// `Arc`. Part numbers are assigned when `put_part` is called, not when its future resolves, so +/// parts stay correctly ordered however they interleave. +#[derive(Debug)] +pub(super) struct MultipartState { + pub(super) store: Arc, + pub(super) location: Path, + pub(super) upload_id: String, + pub(super) parts: Mutex>>, +} + +impl MultipartState { + /// Upload a single part and record its ETag against `part_idx`. + async fn put_part(self: Arc, part_idx: usize, data: PutPayload) -> Result<()> { + let part = self.upload_part(part_idx, data).await?; + let mut parts = self.parts.lock().expect("multipart state poisoned"); + if parts.len() <= part_idx { + parts.resize(part_idx + 1, None); + } + parts[part_idx] = Some(part); + Ok(()) + } + + async fn upload_part(&self, part_idx: usize, data: PutPayload) -> Result { + let mut url = self.store.object_url(&self.location)?; + { + let mut query = url.query_pairs_mut(); + query.append_pair("partNumber", &(part_idx + 1).to_string()); + query.append_pair("uploadId", &self.upload_id); + } + let request = SignedRequest::new(Method::PUT, url) + .header( + CONTENT_LENGTH, + HeaderValue::from_str(&data.content_length().to_string()) + .expect("content length is a valid header value"), + ) + .body(data); + let response = self.store.send(&self.location, request).await?; + let e_tag = header_str(response.headers(), &ETAG).ok_or_else(|| { + error_for( + "UploadPart response is missing an ETag", + std::io::Error::other(format!("part {}", part_idx + 1)), + ) + })?; + Ok(PartId { content_id: e_tag }) + } + + /// The ETags of every part, in order, erroring if `put_part` was not awaited for some part. + fn finished_parts(&self, expected: usize) -> Result> { + let parts = self.parts.lock().expect("multipart state poisoned"); + (0..expected) + .map(|part_idx| { + parts.get(part_idx).and_then(Clone::clone).ok_or_else(|| { + error_for( + "multipart upload completed before all parts finished uploading", + std::io::Error::other(format!("part {} is missing", part_idx + 1)), + ) + }) + }) + .collect() + } +} + +/// A multipart upload in which every request — initiation, each part, completion and abort — is +/// independently signed by the Python callback. +#[derive(Debug)] +pub(super) struct RemoteSignedMultipartUpload { + pub(super) state: Arc, + pub(super) part_idx: usize, +} + +#[async_trait] +impl MultipartUpload for RemoteSignedMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + let part_idx = self.part_idx; + self.part_idx += 1; + let state = Arc::clone(&self.state); + state.put_part(part_idx, data).boxed() + } + + async fn complete(&mut self) -> Result { + let mut parts = self.state.finished_parts(self.part_idx)?; + if parts.is_empty() { + // S3 rejects a completion with no parts, so an empty object still needs one part. + parts.push(self.state.upload_part(0, PutPayload::default()).await?); + self.part_idx = 1; + } + let body = quick_xml::se::to_string(&CompleteMultipartUpload::from(parts)) + .map_err(|error| error_for("failed to encode CompleteMultipartUpload", error))?; + + let mut url = self.state.store.object_url(&self.state.location)?; + url.query_pairs_mut() + .append_pair("uploadId", &self.state.upload_id); + let request = SignedRequest::new(Method::POST, url) + .header( + CONTENT_LENGTH, + HeaderValue::from_str(&body.len().to_string()) + .expect("content length is a valid header value"), + ) + .body(PutPayload::from(body)); + + let response = self.state.store.send(&self.state.location, request).await?; + let version = header_str(response.headers(), &HeaderName::from_static(VERSION_ID)); + let body = + response.into_body().bytes().await.map_err(|error| { + error_for("failed to read CompleteMultipartUpload response", error) + })?; + // S3 can report a failed completion with a 200 status and an `` body, so the body + // has to be inspected rather than trusting the status alone. + if let Some(detail) = s3_error_detail(&body) { + return Err(object_store::Error::Generic { + store: STORE, + source: format!("CompleteMultipartUpload failed: {detail}").into(), + }); + } + let response: CompleteMultipartUploadResult = quick_xml::de::from_reader(body.as_ref()) + .map_err(|error| error_for("invalid CompleteMultipartUpload response", error))?; + Ok(PutResult { + e_tag: Some(response.e_tag), + version, + extensions: Default::default(), + }) + } + + async fn abort(&mut self) -> Result<()> { + let mut url = self.state.store.object_url(&self.state.location)?; + url.query_pairs_mut() + .append_pair("uploadId", &self.state.upload_id); + self.state + .store + .send( + &self.state.location, + SignedRequest::new(Method::DELETE, url), + ) + .await?; + Ok(()) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct InitiateMultipartUploadResult { + upload_id: String, +} + +#[derive(Serialize)] +#[serde(rename_all = "PascalCase")] +struct CompleteMultipartUpload { + part: Vec, +} + +impl From> for CompleteMultipartUpload { + fn from(value: Vec) -> Self { + Self { + part: value + .into_iter() + .enumerate() + .map(|(part_idx, part)| MultipartPart { + e_tag: part.content_id, + part_number: part_idx + 1, + }) + .collect(), + } + } +} + +#[derive(Serialize)] +struct MultipartPart { + #[serde(rename = "ETag")] + e_tag: String, + #[serde(rename = "PartNumber")] + part_number: usize, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct CompleteMultipartUploadResult { + #[serde(rename = "ETag")] + e_tag: String, +} diff --git a/pyo3-object_store/src/aws/remote_signed/signer.rs b/pyo3-object_store/src/aws/remote_signed/signer.rs new file mode 100644 index 00000000..b346763d --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/signer.rs @@ -0,0 +1,135 @@ +//! The Python boundary: handing each request to a user callback to be signed. + +use std::collections::HashMap; +use std::fmt::{Debug, Formatter}; + +use http::{HeaderMap, HeaderName, HeaderValue, Method}; +use object_store::Result; +use pyo3::exceptions::PyTypeError; +use pyo3::intern; +use pyo3::prelude::*; +use url::Url; + +use super::error_for; + +/// A Python callable that signs a single S3 request. +pub(super) struct PySigner(pub(super) Py); + +impl Clone for PySigner { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl Debug for PySigner { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("PySigner(..)") + } +} + +impl<'py> FromPyObject<'_, 'py> for PySigner { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult { + if !obj.hasattr(intern!(obj.py(), "__call__"))? { + return Err(PyTypeError::new_err("Expected callable object for signer.")); + } + Ok(Self(obj.as_unbound().clone_ref(obj.py()))) + } +} + +impl<'py> IntoPyObject<'py> for PySigner { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult> { + Ok(self.0.into_bound(py)) + } +} + +/// A signer callback may be a plain function or a coroutine function. +enum PySignerResult { + Async(Py), + Sync((String, HashMap)), +} + +impl PySignerResult { + async fn resolve(self) -> PyResult<(String, HashMap)> { + match self { + Self::Sync(result) => Ok(result), + Self::Async(coroutine) => { + let future = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.bind(py).clone()) + })?; + let result = future.await?; + Python::attach(|py| result.extract(py)) + } + } + } +} + +impl<'py> FromPyObject<'_, 'py> for PySignerResult { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult { + if obj.hasattr(intern!(obj.py(), "__await__"))? { + Ok(Self::Async(obj.as_unbound().clone_ref(obj.py()))) + } else { + Ok(Self::Sync(obj.extract()?)) + } + } +} + +impl PySigner { + /// Hand `(method, uri, headers)` to the callback and return the signed `(uri, headers)`. + /// + /// The returned headers *replace* the headers passed in, rather than being merged into them. + /// This is what the S3 remote-signing contract specifies: a signer computes a signature over a + /// specific set of headers, so silently adding to or keeping headers it did not return would + /// invalidate that signature. + pub(super) async fn sign( + &self, + method: &Method, + uri: &Url, + headers: &HeaderMap, + ) -> Result<(Url, HeaderMap)> { + let headers: HashMap = headers + .iter() + .map(|(name, value)| { + Ok(( + name.as_str().to_owned(), + value + .to_str() + .map_err(|error| error_for("invalid request header", error))? + .to_owned(), + )) + }) + .collect::>()?; + let result = Python::attach(|py| { + self.0 + .call1(py, (method.as_str(), uri.as_str(), headers))? + .extract::(py) + }) + .map_err(|error| error_for("signer callback failed", error))? + .resolve() + .await + .map_err(|error| error_for("signer callback failed", error))?; + let uri = Url::parse(&result.0) + .map_err(|error| error_for("signer returned invalid URI", error))?; + let headers = result + .1 + .into_iter() + .map(|(name, value)| { + Ok(( + HeaderName::from_bytes(name.as_bytes()) + .map_err(|error| error_for("signer returned invalid header name", error))?, + HeaderValue::from_str(&value).map_err(|error| { + error_for("signer returned invalid header value", error) + })?, + )) + }) + .collect::>()?; + Ok((uri, headers)) + } +} diff --git a/pyo3-object_store/src/aws/remote_signed/store.rs b/pyo3-object_store/src/aws/remote_signed/store.rs new file mode 100644 index 00000000..fe7cefa4 --- /dev/null +++ b/pyo3-object_store/src/aws/remote_signed/store.rs @@ -0,0 +1,305 @@ +//! The Python-facing class: construction, configuration, pickling. + +use std::sync::Arc; + +use object_store::client::{HttpConnector, ReqwestConnector}; +use object_store::path::Path; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; +use url::Url; + +use crate::error::PyObjectStoreResult; +use crate::prefix::MaybePrefixedStore; +use crate::retry::PyRetryConfig; +use crate::{PyClientOptions, PyUrl}; + +use super::list::decode_prefix; +use super::signer::PySigner; +use super::{RemoteSignedS3Store, STORE}; + +/// Resolve an `s3://bucket/prefix` location against an S3 endpoint into the HTTPS URL that +/// [`PyRemoteSignedS3Store::new`] takes. +fn s3_endpoint_url( + url: &Url, + endpoint: &Url, + virtual_hosted_style_request: bool, +) -> PyResult { + if !matches!(url.scheme(), "s3" | "s3a") { + return Err(PyValueError::new_err(format!( + "Expected an s3:// or s3a:// URL, got {}. Pass an HTTPS endpoint URL to \ + RemoteSignedS3Store() directly instead.", + url.scheme(), + ))); + } + let bucket = url + .host_str() + .filter(|bucket| !bucket.is_empty()) + .ok_or_else(|| PyValueError::new_err(format!("{url} does not name a bucket")))?; + let prefix = url + .path_segments() + .into_iter() + .flatten() + .filter(|segment| !segment.is_empty()) + .collect::>() + .join("/"); + + if !matches!(endpoint.scheme(), "http" | "https") { + return Err(PyValueError::new_err(format!( + "endpoint must be an http:// or https:// URL, got {endpoint}", + ))); + } + // An endpoint served below a path would make the split between endpoint, bucket and key + // prefix ambiguous, so require a bare origin and let the caller assemble the URL itself. + if !matches!(endpoint.path(), "" | "/") { + return Err(PyValueError::new_err(format!( + "endpoint must not include a path, got {}. Build the full URL yourself and pass \ + it to RemoteSignedS3Store() instead.", + endpoint.path(), + ))); + } + + let mut resolved = endpoint.clone(); + resolved.set_query(None); + resolved.set_fragment(None); + if virtual_hosted_style_request { + let host = endpoint.host_str().ok_or_else(|| { + PyValueError::new_err(format!("endpoint {endpoint} does not include a host")) + })?; + resolved + .set_host(Some(&format!("{bucket}.{host}"))) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + resolved.set_path(&format!("/{prefix}")); + } else { + resolved.set_path(&format!("/{bucket}/{prefix}")); + } + Ok(resolved.into()) +} + +/// The constructor arguments of a [`PyRemoteSignedS3Store`], retained for pickling. +#[derive(Debug, Clone)] +struct RemoteSignedS3Config { + url: String, + signer: PySigner, + virtual_hosted_style_request: bool, + client_options: Option, + retry_config: Option, + /// The bucket named by `url`, and the key prefix it implies. Derived rather than passed, + /// but kept here so the Python getters can report them. + bucket: String, + prefix: Option, +} + +impl RemoteSignedS3Config { + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = + PyTuple::new(py, [self.url.clone().into_bound_py_any(py)?])?.into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "signer"), self.signer.clone())?; + if self.virtual_hosted_style_request { + kwargs.set_item(intern!(py, "virtual_hosted_style_request"), true)?; + } + if let Some(client_options) = &self.client_options { + kwargs.set_item(intern!(py, "client_options"), client_options.clone())?; + } + if let Some(retry_config) = &self.retry_config { + kwargs.set_item(intern!(py, "retry_config"), retry_config.clone())?; + } + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } + + /// Whether two stores were built from equivalent arguments. + /// + /// The signer is compared by identity: two distinct callables cannot be shown to sign + /// equivalently. + fn eq(&self, other: &Self, py: Python<'_>) -> bool { + self.url == other.url + && self.virtual_hosted_style_request == other.virtual_hosted_style_request + && self.client_options == other.client_options + && self.retry_config == other.retry_config + && self.signer.0.bind(py).is(other.signer.0.bind(py)) + } +} + +/// Python-native wrapper for [`RemoteSignedS3Store`]. +#[derive(Debug, Clone)] +#[pyclass(name = "RemoteSignedS3Store", frozen, subclass, from_py_object)] +pub struct PyRemoteSignedS3Store { + store: Arc>, + /// The arguments used for pickling. This must stay in sync with the underlying store. + config: RemoteSignedS3Config, +} + +impl AsRef>> for PyRemoteSignedS3Store { + fn as_ref(&self) -> &Arc> { + &self.store + } +} + +impl PyRemoteSignedS3Store { + /// Consume self and return the underlying [`RemoteSignedS3Store`]. + pub fn into_inner(self) -> Arc> { + self.store + } +} + +#[pymethods] +impl PyRemoteSignedS3Store { + #[new] + #[pyo3(signature = (url, signer, *, virtual_hosted_style_request=false, client_options=None, retry_config=None))] + fn new( + url: String, + signer: PySigner, + virtual_hosted_style_request: bool, + client_options: Option, + retry_config: Option, + ) -> PyObjectStoreResult { + let parsed = Url::parse(&url).map_err(|error| PyValueError::new_err(error.to_string()))?; + let mut segments = parsed + .path_segments() + .ok_or_else(|| PyValueError::new_err(format!("{STORE} URL must include a path")))? + .filter(|segment| !segment.is_empty()); + + // In path style the first segment names the bucket and the rest is the key prefix; in + // virtual-hosted style the bucket is the host's first label, so every segment is prefix. + let mut bucket_url = parsed.clone(); + bucket_url.set_query(None); + bucket_url.set_fragment(None); + let (bucket, prefix) = if virtual_hosted_style_request { + let host = parsed.host_str().ok_or_else(|| { + PyValueError::new_err(format!( + "{STORE} URL must include a host naming the bucket when \ + virtual_hosted_style_request is set" + )) + })?; + let bucket = host.split('.').next().unwrap_or(host).to_owned(); + bucket_url.set_path("/"); + (bucket, decode_prefix(segments)) + } else { + // The bucket segment needs no decoding: S3 bucket names are restricted to + // characters that a URL never percent-encodes. + let bucket = segments + .next() + .ok_or_else(|| PyValueError::new_err(format!("{STORE} URL must include a bucket")))? + .to_owned(); + bucket_url.set_path(&format!("/{bucket}/")); + (bucket, decode_prefix(segments)) + }; + // An empty prefix stays `None`, so that `MaybePrefixedStore` passes paths straight + // through instead of rewriting them. + let prefix = if prefix.is_empty() { + None + } else { + Some(Path::parse(prefix).map_err(|error| PyValueError::new_err(error.to_string()))?) + }; + + let options = client_options.clone().map(Into::into).unwrap_or_default(); + let client = ReqwestConnector {}.connect(&options)?; + let store = RemoteSignedS3Store { + bucket_url, + bucket: bucket.clone(), + signer: signer.clone(), + client, + retry: retry_config.clone().map(Into::into).unwrap_or_default(), + }; + Ok(Self { + store: Arc::new(MaybePrefixedStore::new(store, prefix.clone())), + config: RemoteSignedS3Config { + url, + signer, + virtual_hosted_style_request, + client_options, + retry_config, + bucket, + prefix, + }, + }) + } + + /// Construct a store from an `s3://` location plus the S3 endpoint that serves it. + /// + /// Catalogs hand out locations as `s3://bucket/key` and the endpoint separately, so this + /// saves callers assembling the HTTPS URL themselves. + #[classmethod] + #[pyo3(signature = (url, signer, *, endpoint, virtual_hosted_style_request=false, client_options=None, retry_config=None))] + fn from_s3_url<'py>( + cls: &Bound<'py, PyType>, + url: PyUrl, + signer: PySigner, + endpoint: PyUrl, + virtual_hosted_style_request: bool, + client_options: Option, + retry_config: Option, + ) -> PyObjectStoreResult> { + let url = s3_endpoint_url( + url.as_ref(), + endpoint.as_ref(), + virtual_hosted_style_request, + )?; + + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = PyDict::new(cls.py()); + kwargs.set_item(intern!(cls.py(), "signer"), signer)?; + kwargs.set_item( + intern!(cls.py(), "virtual_hosted_style_request"), + virtual_hosted_style_request, + )?; + kwargs.set_item(intern!(cls.py(), "client_options"), client_options)?; + kwargs.set_item(intern!(cls.py(), "retry_config"), retry_config)?; + Ok(cls.call((url,), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not the same + // type + other + .cast::() + .map(|other| self.config.eq(&other.get().config, other.py())) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + format!("RemoteSignedS3Store(\"{}\")", self.config.url) + } + + #[getter] + fn url(&self) -> &str { + &self.config.url + } + + #[getter] + fn bucket(&self) -> &str { + &self.config.bucket + } + + #[getter] + fn prefix(&self) -> Option<&str> { + self.config.prefix.as_ref().map(Path::as_ref) + } + + #[getter] + fn signer(&self) -> PySigner { + self.config.signer.clone() + } + + #[getter] + fn virtual_hosted_style_request(&self) -> bool { + self.config.virtual_hosted_style_request + } + + #[getter] + fn client_options(&self) -> Option { + self.config.client_options.clone() + } + + #[getter] + fn retry_config(&self) -> Option { + self.config.retry_config.clone() + } +} diff --git a/pyo3-object_store/src/lib.rs b/pyo3-object_store/src/lib.rs index 0d7d7fdf..1147d046 100644 --- a/pyo3-object_store/src/lib.rs +++ b/pyo3-object_store/src/lib.rs @@ -20,7 +20,7 @@ mod store; mod url; pub use api::{register_exceptions_module, register_store_module}; -pub use aws::PyS3Store; +pub use aws::{PyRemoteSignedS3Store, PyS3Store}; pub use azure::PyAzureStore; pub use client::{PyClientConfigKey, PyClientOptions}; pub use error::{PyObjectStoreError, PyObjectStoreResult}; diff --git a/pyo3-object_store/src/store.rs b/pyo3-object_store/src/store.rs index 2b802f1e..dc68a37b 100644 --- a/pyo3-object_store/src/store.rs +++ b/pyo3-object_store/src/store.rs @@ -7,7 +7,10 @@ use pyo3::pybacked::PyBackedStr; use pyo3::types::{PyDict, PyTuple}; use pyo3::{intern, PyTypeInfo}; -use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store}; +use crate::{ + PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyRemoteSignedS3Store, + PyS3Store, +}; /// A wrapper around a Rust ObjectStore instance that allows any rust-native implementation of /// ObjectStore. @@ -33,6 +36,8 @@ impl<'py> FromPyObject<'_, 'py> for PyObjectStore { Ok(Self(store.get().as_ref().clone())) } else if let Ok(store) = obj.cast::() { Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) } else { let py = obj.py(); // Check for object-store instance from other library @@ -46,6 +51,7 @@ impl<'py> FromPyObject<'_, 'py> for PyObjectStore { PyHttpStore::type_object(py).name()?.to_str()?, PyLocalStore::type_object(py).name()?.to_str()?, PyMemoryStore::type_object(py).name()?.to_str()?, + PyRemoteSignedS3Store::type_object(py).name()?.to_str()?, PyS3Store::type_object(py).name()?.to_str()?, ] .contains(&cls_name.as_str()) @@ -147,6 +153,18 @@ impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStoreInner { return Ok(Self(store.into_inner())); } + if cls_name.as_str() == PyRemoteSignedS3Store::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyRemoteSignedS3Store::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + if cls_name.as_str() == PyS3Store::type_object(py).name()? { let (args, kwargs): (Bound, Bound) = obj .call_method0(intern!(py, "__getnewargs_ex__"))? diff --git a/tests/store/fake_s3.py b/tests/store/fake_s3.py new file mode 100644 index 00000000..4dc81d27 --- /dev/null +++ b/tests/store/fake_s3.py @@ -0,0 +1,416 @@ +"""An in-process S3 that rejects unsigned requests. + +Used by the `RemoteSignedS3Store` tests. Only the subset of the S3 REST API that the +store issues is implemented, and every response is deliberately close to what a real S3 +returns — including its `` documents and its opaque, `+`-bearing upload ids — +because the store's correctness depends on those details. +""" + +from __future__ import annotations + +import re +from contextlib import contextmanager +from datetime import timedelta +from email.utils import parsedate_to_datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Lock, Thread +from typing import TYPE_CHECKING, ClassVar, NamedTuple +from urllib.parse import parse_qs, quote, unquote, urlsplit + +from obstore.store import RemoteSignedS3Store + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from obstore.store import ClientConfig, RetryConfig + +DATA = b"abcdefghijklmnopqrstuvwxyz" + +LAST_MODIFIED = "Wed, 21 Oct 2015 07:28:00 GMT" +"""Every object reports this mtime, so date preconditions are predictable.""" + +LAST_MODIFIED_AT = parsedate_to_datetime(LAST_MODIFIED) + +STORED_HEADERS = frozenset( + { + "cache-control", + "content-disposition", + "content-encoding", + "content-language", + "content-type", + "x-amz-storage-class", + }, +) +"""Headers a real S3 stores with the object and echoes back on GET.""" + + +class Request(NamedTuple): + """One request that reached the server, recorded for assertions.""" + + method: str + path: str + headers: dict[str, str] + + +CLIENT_OPTIONS: ClientConfig = {"allow_http": True} +"""The test server speaks plain HTTP, which `allow_http` must opt into.""" + +NO_WAIT_RETRIES: RetryConfig = { + "max_retries": 3, + "backoff": {"init_backoff": timedelta(0), "max_backoff": timedelta(0), "base": 1}, +} +"""Retry immediately, so retry tests do not spend real time sleeping.""" + + +class FakeS3(BaseHTTPRequestHandler): + """A minimal S3 that rejects unsigned requests. + + Only the subset of the S3 REST API that `RemoteSignedS3Store` issues is implemented. + State is class-level because `ThreadingHTTPServer` builds a handler per request. + """ + + lock: ClassVar[Lock] = Lock() + objects: ClassVar[dict[str, bytes]] = {} + attributes: ClassVar[dict[str, dict[str, str]]] = {} + """Per-object headers a real S3 would store and echo back on GET.""" + versions: ClassVar[dict[str, dict[str, bytes]]] = {} + """Superseded object versions, keyed by object key then version id.""" + uploads: ClassVar[dict[str, dict[int, bytes]]] = {} + page_size: ClassVar[int] = 1000 + fail_next: ClassVar[int] = 0 + """Number of upcoming requests to fail with a retryable 503.""" + signed_requests: ClassVar[list[Request]] = [] + """Every request that passed signature checking, in arrival order.""" + + @classmethod + def reset(cls) -> None: + """Forget all state from a previous test.""" + cls.objects = {} + cls.attributes = {} + cls.versions = {} + cls.uploads = {} + cls.page_size = 1000 + cls.fail_next = 0 + cls.signed_requests = [] + + # ---- request plumbing ---- + + @property + def key(self) -> str: + """The object key addressed by this request, with the bucket stripped.""" + path = unquote(urlsplit(self.path).path) + return path.removeprefix("/bucket").lstrip("/") + + @property + def query(self) -> dict[str, list[str]]: + """The parsed query string, keeping valueless keys such as `?uploads`.""" + return parse_qs(urlsplit(self.path).query, keep_blank_values=True) + + def read_body(self) -> bytes: + """Read exactly the body the client announced.""" + return self.rfile.read(int(self.headers.get("Content-Length", 0))) + + def error(self, status: int, code: str, message: str) -> None: + """Reply with an S3 `` document, as a real S3 would.""" + body = ( + f'{code}' + f"{message}" + ).encode() + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def reply(self, status: int, body: bytes = b"", **headers: str) -> None: + """Reply with `body`, omitting it for a HEAD as HTTP requires.""" + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + for name, value in headers.items(): + self.send_header(name.replace("_", "-"), value) + self.end_headers() + if self.command != "HEAD": + self.wfile.write(body) + + def stored_attributes(self, key: str) -> dict[str, str]: + """Return `key`'s attribute headers, in `reply()`'s underscore form.""" + with FakeS3.lock: + stored = FakeS3.attributes.get(key, {}) + return {name.replace("-", "_"): value for name, value in stored.items()} + + def precondition_failure(self) -> bool: + """Answer any date precondition the fixed `Last-Modified` fails to satisfy.""" + since = self.headers.get("if-modified-since") + if since and parsedate_to_datetime(since) >= LAST_MODIFIED_AT: + self.reply(304) + return True + since = self.headers.get("if-unmodified-since") + if since and parsedate_to_datetime(since) < LAST_MODIFIED_AT: + self.error(412, "PreconditionFailed", "Object modified since then") + return True + return False + + def authorize(self) -> bool: + """Reject any request the signer did not sign, and fail injected requests.""" + if self.headers.get("x-signed") != "yes": + self.error(403, "AccessDenied", "Request was not signed") + return False + with FakeS3.lock: + FakeS3.signed_requests.append( + Request(self.command, self.path, dict(self.headers.items())), + ) + if FakeS3.fail_next: + FakeS3.fail_next -= 1 + self.error(503, "SlowDown", "Please retry") + return False + return True + + # ---- verbs ---- + + def do_GET(self) -> None: + """Serve a list, or a whole or ranged object.""" + if not self.authorize(): + return + if self.query.get("list-type") == ["2"]: + self.serve_list() + return + + if version := self.query.get("versionId"): + with FakeS3.lock: + body = FakeS3.versions.get(self.key, {}).get(version[0]) + if body is None: + self.error(404, "NoSuchVersion", f"No version {version[0]}") + return + else: + with FakeS3.lock: + body = FakeS3.objects.get(self.key) + if body is None: + self.error(404, "NoSuchKey", "The specified key does not exist") + return + + if self.precondition_failure(): + return + + common = { + "Last_Modified": LAST_MODIFIED, + "ETag": '"etag"', + **self.stored_attributes(self.key), + } + if range_header := self.headers.get("range"): + start, _, end = range_header.removeprefix("bytes=").partition("-") + start, end = int(start), int(end) + 1 + self.reply( + 206, + body[start:end], + Content_Range=f"bytes {start}-{end - 1}/{len(body)}", + **common, + ) + else: + self.reply(200, body, **common) + + def do_HEAD(self) -> None: + """Serve object metadata.""" + self.do_GET() + + def do_PUT(self) -> None: + """Store a whole object or one part of a multipart upload.""" + if not self.authorize(): + return + body = self.read_body() + query = self.query + + if upload_id := query.get("uploadId"): + part_number = int(query["partNumber"][0]) + with FakeS3.lock: + FakeS3.uploads[upload_id[0]][part_number] = body + self.reply(200, ETag=f'"part-{part_number}"') + return + + if source := self.headers.get("x-amz-copy-source"): + # A real copy takes the object and its attributes from the source key. + source_key = unquote(source).removeprefix("/bucket").lstrip("/") + with FakeS3.lock: + if source_key not in FakeS3.objects: + self.error(404, "NoSuchKey", f"No such source {source_key}") + return + body = FakeS3.objects[source_key] + attributes = dict(FakeS3.attributes.get(source_key, {})) + else: + attributes = { + name: value + for name, value in self.headers.items() + if name.lower() in STORED_HEADERS + or name.lower().startswith("x-amz-meta-") + } + + with FakeS3.lock: + exists = self.key in FakeS3.objects + if self.headers.get("if-none-match") == "*" and exists: + self.error(412, "PreconditionFailed", "Object already exists") + return + if exists: + # Keep the superseded bytes addressable by version id. + previous = FakeS3.versions.setdefault(self.key, {}) + previous[f"v{len(previous)}"] = FakeS3.objects[self.key] + FakeS3.objects[self.key] = body + FakeS3.attributes[self.key] = attributes + version = f"v{len(FakeS3.versions.get(self.key, {}))}" + self.reply(200, ETag='"etag"', x_amz_version_id=version) + + def do_POST(self) -> None: + """Initiate or complete a multipart upload.""" + if not self.authorize(): + return + body = self.read_body() + query = self.query + + if "uploads" in query: + # Real S3 upload ids are opaque base64-ish strings, so they routinely + # contain `+`, `/` and `=`. Using one here keeps the client honest about + # encoding them: sent raw in a query string, `+` arrives as a space and + # then addresses no upload at all. + upload_id = f"up+{len(FakeS3.uploads)}/id==" + with FakeS3.lock: + FakeS3.uploads[upload_id] = {} + self.reply( + 200, + ( + '' + f"{upload_id}" + "" + ).encode(), + ) + return + + upload_id = query["uploadId"][0] + if upload_id not in FakeS3.uploads: + # Answer like S3 rather than raising: an unhandled exception closes the + # connection, which the client treats as retryable and then backs off for + # minutes, turning a regression into a hang instead of a failure. + self.error(404, "NoSuchUpload", f"Unknown upload {upload_id}") + return + # Assemble the object from the parts the client listed, in the order it listed + # them, so that a wrongly ordered CompleteMultipartUpload produces wrong bytes. + part_numbers = [ + int(match) for match in re.findall(r"(\d+)<", body.decode()) + ] + with FakeS3.lock: + parts = FakeS3.uploads.pop(upload_id) + FakeS3.objects[self.key] = b"".join( + parts[number] for number in part_numbers + ) + self.reply( + 200, + ( + b'' + b""multipart-etag"" + b"" + ), + ) + + def do_DELETE(self) -> None: + """Delete an object or abort a multipart upload.""" + if not self.authorize(): + return + with FakeS3.lock: + if upload_id := self.query.get("uploadId"): + FakeS3.uploads.pop(upload_id[0], None) + else: + FakeS3.objects.pop(self.key, None) + self.reply(204) + + def serve_list(self) -> None: + """Serve one page of a `ListObjectsV2` response.""" + query = self.query + prefix = query.get("prefix", [""])[0] + delimiter = query.get("delimiter", [None])[0] + start_after = query.get("continuation-token", query.get("start-after", [""]))[0] + + with FakeS3.lock: + keys = sorted(key for key in FakeS3.objects if key.startswith(prefix)) + keys = [key for key in keys if key > start_after] + + contents, prefixes = [], [] + for key in keys: + tail = key[len(prefix) :] + if delimiter and delimiter in tail: + prefixes.append(prefix + tail.split(delimiter)[0] + delimiter) + else: + contents.append(key) + + truncated = len(contents) > FakeS3.page_size + page = contents[: FakeS3.page_size] + body = [''] + for key in page: + with FakeS3.lock: + size = len(FakeS3.objects[key]) + body.append( + f"{quote(key)}" + "2015-10-21T07:28:00.000Z" + f"{size}"etag"", + ) + body.extend( + f"{quote(value)}" + for value in dict.fromkeys(prefixes) + ) + if truncated: + body.append(f"{page[-1]}") + body.append("") + self.reply(200, "".join(body).encode()) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 + """Keep the test server quiet.""" + + +@contextmanager +def signed_server() -> Iterator[str]: + """Run `FakeS3` on a free port and yield its bucket URL.""" + FakeS3.reset() + server = ThreadingHTTPServer(("127.0.0.1", 0), FakeS3) + thread = Thread(target=server.serve_forever) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/bucket" + finally: + server.shutdown() + thread.join() + + +def signer( + _method: str, + uri: str, + headers: dict[str, str], +) -> tuple[str, dict[str, str]]: + """Sign every request by stamping it with the header `FakeS3` demands.""" + return uri, {**headers, "x-signed": "yes"} + + +def make_store( + url: str, + *, + virtual_hosted_style_request: bool = False, + retry_config: RetryConfig | None = None, +) -> RemoteSignedS3Store: + """Build a store against the test server.""" + return RemoteSignedS3Store( + url, + signer, + virtual_hosted_style_request=virtual_hosted_style_request, + client_options=CLIENT_OPTIONS, + retry_config=retry_config, + ) + + +def select(method: str, query_key: str) -> list[Request]: + """Return the signed requests using `method` whose query mentions `query_key`.""" + return [ + request + for request in FakeS3.signed_requests + if request.method == method and query_key in urlsplit(request.path).query + ] + + +def one(requests: Iterable[Request]) -> Request: + """Return the single matching request, asserting that there is exactly one.""" + matched = list(requests) + assert len(matched) == 1, f"expected exactly one request, got {len(matched)}" + return matched[0] diff --git a/tests/store/test_remote_signed_s3.py b/tests/store/test_remote_signed_s3.py new file mode 100644 index 00000000..7ba56d58 --- /dev/null +++ b/tests/store/test_remote_signed_s3.py @@ -0,0 +1,630 @@ +"""Tests for `RemoteSignedS3Store`, against the signing-aware server in `fake_s3`.""" + +from __future__ import annotations + +import asyncio +import pickle +from datetime import timedelta +from typing import TYPE_CHECKING +from urllib.parse import urlsplit + +import pytest +import requests + +import obstore as obs +from obstore.exceptions import ( + AlreadyExistsError, + GenericError, + NotModifiedError, + PermissionDeniedError, + PreconditionError, +) +from obstore.store import RemoteSignedS3Store + +from .fake_s3 import ( + CLIENT_OPTIONS, + DATA, + LAST_MODIFIED_AT, + NO_WAIT_RETRIES, + FakeS3, + make_store, + one, + select, + signed_server, + signer, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + +pytestmark = pytest.mark.thread_unsafe( + reason=( + "FakeS3 keeps its state — objects, signed_requests, uploads — in " + "class-level attributes shared by every server instance. Running a " + "test body twice concurrently, as --parallel-threads does, races on " + "that shared state instead of exercising RemoteSignedS3Store." + ), +) + + +@pytest.mark.asyncio +async def test_signed_range_request() -> None: + with signed_server() as url: + response = await asyncio.to_thread(requests.get, f"{url}/chunk", timeout=1) + assert response.status_code == 403 + + calls: list[tuple[str, str, dict[str, str]]] = [] + + async def recording_signer(method: str, uri: str, headers: dict[str, str]): + calls.append((method, uri, headers)) + return signer(method, uri, headers) + + store = RemoteSignedS3Store( + url, + recording_signer, + client_options=CLIENT_OPTIONS, + ) + FakeS3.objects["chunk"] = DATA + + assert await store.get_range_async("chunk", start=5, end=10) == DATA[5:10] + assert (await store.head_async("chunk"))["size"] == len(DATA) + + assert calls == [ + ("GET", f"{url}/chunk", {"range": "bytes=5-9"}), + ("HEAD", f"{url}/chunk", {}), + ] + + +def test_unsigned_request_is_rejected_with_server_message() -> None: + with signed_server() as url: + store = RemoteSignedS3Store( + url, + lambda _method, uri, headers: (uri, headers), + client_options=CLIENT_OPTIONS, + ) + with pytest.raises( + PermissionDeniedError, + match="AccessDenied: Request was not signed", + ): + store.get("chunk").bytes() + + +def test_put_and_get_roundtrip() -> None: + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + assert store.get("chunk").bytes() == DATA + + +def test_get_missing_object_raises_not_found() -> None: + with signed_server() as url: + store = make_store(url) + with pytest.raises(FileNotFoundError, match="NoSuchKey"): + store.get("missing").bytes() + + +def test_key_with_reserved_characters_is_encoded() -> None: + """A `?` in a key must be percent-encoded, not read as the start of a query.""" + with signed_server() as url: + store = make_store(url) + store.put("a b?c#d/chunk", DATA) + assert store.get("a b?c#d/chunk").bytes() == DATA + assert [meta["path"] for meta in store.list().collect()] == ["a b?c#d/chunk"] + + +@pytest.mark.parametrize( + ("encoded", "decoded"), + [ + ("pre%20fix", "pre fix"), + ("caf%C3%A9", "café"), + ("a%23b", "a#b"), + ("p%25c", "p%c"), + ], +) +def test_url_prefix_is_percent_decoded(encoded: str, decoded: str) -> None: + """A URL can only express these prefixes encoded; keys must come out decoded. + + Taking the segment literally would double-encode it and address a different object. + """ + with signed_server() as url: + store = make_store(f"{url}/{encoded}") + assert store.prefix == decoded + store.put("chunk", DATA) + assert store.get("chunk").bytes() == DATA + assert [meta["path"] for meta in store.list().collect()] == ["chunk"] + + assert FakeS3.objects.keys() == {f"{decoded}/chunk"} + + +def test_signed_list_paginates() -> None: + with signed_server() as url: + FakeS3.page_size = 2 + store = make_store(f"{url}/prefix") + for index in range(5): + store.put(f"chunk-{index}", DATA) + + objects = store.list().collect() + + assert [meta["path"] for meta in objects] == [ + f"chunk-{index}" for index in range(5) + ] + assert objects[0]["size"] == len(DATA) + list_requests = [ + request.path + for request in FakeS3.signed_requests + if "list-type" in request.path + ] + assert len(list_requests) == 3, "expected three pages" + assert "prefix=prefix%2F" in list_requests[0] + + +def test_list_prefix_matches_on_segment_boundaries() -> None: + """Listing `a/b` must not also return `a/bc/...`. + + The listed prefix is sent with a trailing delimiter for exactly this reason. + """ + with signed_server() as url: + store = make_store(f"{url}/root") + store.put("a/b/inside", DATA) + store.put("a/bc/outside", DATA) + + assert [meta["path"] for meta in store.list("a/b").collect()] == ["a/b/inside"] + + assert FakeS3.objects.keys() == {"root/a/b/inside", "root/a/bc/outside"} + + +def test_list_with_delimiter_returns_common_prefixes() -> None: + with signed_server() as url: + store = make_store(f"{url}/prefix") + store.put("top", DATA) + store.put("nested/inner", DATA) + + result = store.list_with_delimiter() + + assert [meta["path"] for meta in result["objects"]] == ["top"] + assert [str(prefix) for prefix in result["common_prefixes"]] == ["nested"] + + +def test_conditional_create_conflict_raises_already_exists() -> None: + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA, mode="create") + with pytest.raises(AlreadyExistsError): + store.put("chunk", DATA, mode="create") + + +def test_delete_removes_object() -> None: + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + store.delete("chunk") + assert store.list().collect() == [] + + +def test_get_with_version_reads_the_superseded_object() -> None: + """`GetOptions.version` must reach S3 as `?versionId=`, not be silently ignored.""" + with signed_server() as url: + store = make_store(url) + first = store.put("chunk", b"first") + store.put("chunk", b"second") + + assert store.get("chunk").bytes() == b"second" + version = first["version"] + assert version is not None, "PUT should report x-amz-version-id" + assert store.get("chunk", options={"version": version}).bytes() == b"first" + + assert "versionId=" in one(select("GET", "versionId")).path + + +def test_get_with_unknown_version_is_not_found() -> None: + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + with pytest.raises(FileNotFoundError, match="NoSuchVersion"): + store.get("chunk", options={"version": "v99"}).bytes() + + +def test_if_modified_since_is_sent_as_an_http_date() -> None: + """The store must format the datetime as an HTTP date the server can parse.""" + after = LAST_MODIFIED_AT + timedelta(days=1) + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + with pytest.raises(NotModifiedError): + store.get("chunk", options={"if_modified_since": after}).bytes() + + sent = one(request for request in FakeS3.signed_requests if request.method == "GET") + assert sent.headers["if-modified-since"] == "Thu, 22 Oct 2015 07:28:00 GMT" + + +def test_if_unmodified_since_rejects_a_newer_object() -> None: + before = LAST_MODIFIED_AT - timedelta(days=1) + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + with pytest.raises(PreconditionError, match="PreconditionFailed"): + store.get("chunk", options={"if_unmodified_since": before}).bytes() + + +def test_attributes_survive_a_put_and_get_roundtrip() -> None: + """What `put` stores as headers, `get` must recover as attributes.""" + sent = { + "Content-Type": "application/json", + "Content-Encoding": "gzip", + "Cache-Control": "max-age=60", + } + with signed_server() as url: + store = make_store(url) + obs.put(store, "chunk", DATA, attributes=sent) + + got = dict(store.get("chunk").attributes) + + assert {key: got[key] for key in sent} == sent + + +def test_user_metadata_survives_a_roundtrip() -> None: + with signed_server() as url: + store = make_store(url) + obs.put(store, "chunk", DATA, attributes={"owner": "nathan"}) + + assert dict(store.get("chunk").attributes)["owner"] == "nathan" + + put = one(request for request in FakeS3.signed_requests if request.method == "PUT") + assert put.headers["x-amz-meta-owner"] == "nathan" + + +def test_list_with_offset_is_pushed_down_to_start_after() -> None: + """The offset must reach S3, not be applied client-side after listing all keys.""" + with signed_server() as url: + store = make_store(f"{url}/root") + for index in range(5): + store.put(f"c-{index}", DATA) + + got = [meta["path"] for meta in store.list(offset="c-2").collect()] + + assert got == ["c-3", "c-4"] + assert "start-after=root%2Fc-2" in one(select("GET", "start-after")).path + + +def test_copy_copies_the_bytes() -> None: + with signed_server() as url: + store = make_store(url) + store.put("source", b"payload") + store.copy("source", "target") + + assert store.get("target").bytes() == b"payload" + + +def test_copy_if_not_exists_refuses_to_overwrite() -> None: + with signed_server() as url: + store = make_store(url) + store.put("source", b"payload") + store.put("target", b"existing") + + with pytest.raises(AlreadyExistsError): + obs.copy(store, "source", "target", overwrite=False) + + assert store.get("target").bytes() == b"existing" + assert ( + one( + request + for request in FakeS3.signed_requests + if "x-amz-copy-source" in request.headers + ).headers["if-none-match"] + == "*" + ) + + +def test_rename_moves_the_object() -> None: + """`rename` has no S3 primitive; it must fall back to copy-then-delete.""" + with signed_server() as url: + store = make_store(url) + store.put("source", b"payload") + store.rename("source", "target") + + assert store.get("target").bytes() == b"payload" + assert [meta["path"] for meta in store.list().collect()] == ["target"] + + +def test_copy_addresses_the_source_by_bucket_and_key() -> None: + with signed_server() as url: + store = make_store(f"{url}/prefix") + store.put("source", DATA) + store.copy("source", "target") + + copy = one( + request + for request in FakeS3.signed_requests + if "x-amz-copy-source" in request.headers + ) + assert urlsplit(copy.path).path == "/bucket/prefix/target" + assert copy.headers["x-amz-copy-source"] == "/bucket/prefix/source" + + +@pytest.mark.parametrize( + ("location", "endpoint", "vhost", "expected_url", "expected_prefix"), + [ + ( + "s3://warehouse/zarr/array", + "https://s3.eu-west-1.amazonaws.com", + False, + "https://s3.eu-west-1.amazonaws.com/warehouse/zarr/array", + "zarr/array", + ), + # A bucket-only location has no prefix. + ( + "s3://warehouse", + "http://minio:9000", + False, + "http://minio:9000/warehouse/", + None, + ), + # `s3a://` is accepted, as Hadoop-style catalogs emit it. + ( + "s3a://warehouse/zarr", + "http://minio:9000/", + False, + "http://minio:9000/warehouse/zarr", + "zarr", + ), + # Virtual-hosted style puts the bucket in the host instead. + ( + "s3://warehouse/zarr", + "https://s3.eu-west-1.amazonaws.com", + True, + "https://warehouse.s3.eu-west-1.amazonaws.com/zarr", + "zarr", + ), + ], +) +def test_from_s3_url_resolves_the_location_against_the_endpoint( + location: str, + endpoint: str, + vhost: bool, # noqa: FBT001 + expected_url: str, + expected_prefix: str | None, +) -> None: + store = RemoteSignedS3Store.from_s3_url( + location, + signer, + endpoint=endpoint, + virtual_hosted_style_request=vhost, + ) + assert store.url == expected_url + assert store.bucket == "warehouse" + assert store.prefix == expected_prefix + + +def test_from_s3_url_rejects_a_non_s3_location() -> None: + with pytest.raises(ValueError, match="Expected an s3:// or s3a:// URL"): + RemoteSignedS3Store.from_s3_url( + "https://s3.example.com/warehouse", + signer, + endpoint="https://s3.example.com", + ) + + +def test_from_s3_url_rejects_an_endpoint_with_a_path() -> None: + """A path on the endpoint would make the bucket/prefix split ambiguous.""" + with pytest.raises(ValueError, match="endpoint must not include a path"): + RemoteSignedS3Store.from_s3_url( + "s3://warehouse/zarr", + signer, + endpoint="https://gateway.example.com/s3", + ) + + +def test_from_s3_url_rejects_a_non_http_endpoint() -> None: + with pytest.raises(ValueError, match="endpoint must be an http"): + RemoteSignedS3Store.from_s3_url( + "s3://warehouse/zarr", + signer, + endpoint="s3://warehouse", + ) + + +def test_from_s3_url_reads_and_writes() -> None: + """A store from an `s3://` location addresses the same keys as the plain one.""" + with signed_server() as url: + origin, _, bucket = url.rpartition("/") + store = RemoteSignedS3Store.from_s3_url( + f"s3://{bucket}/prefix", + signer, + endpoint=origin, + client_options=CLIENT_OPTIONS, + ) + store.put("chunk", DATA) + assert store.get("chunk").bytes() == DATA + assert [meta["path"] for meta in store.list().collect()] == ["chunk"] + + assert FakeS3.objects == {"prefix/chunk": DATA} + + +def test_from_s3_url_store_is_picklable() -> None: + with signed_server() as url: + origin, _, bucket = url.rpartition("/") + store = RemoteSignedS3Store.from_s3_url( + f"s3://{bucket}/prefix", + signer, + endpoint=origin, + client_options=CLIENT_OPTIONS, + ) + store.put("chunk", DATA) + + restored = pickle.loads(pickle.dumps(store)) + assert restored == store + assert restored.get("chunk").bytes() == DATA + + +def test_virtual_hosted_style_takes_the_bucket_from_the_host() -> None: + """`x-amz-copy-source` names the bucket even when the request URL never does.""" + store = RemoteSignedS3Store( + "https://mybucket.s3.eu-west-1.amazonaws.com/prefix", + signer, + virtual_hosted_style_request=True, + ) + assert store.bucket == "mybucket" + assert store.prefix == "prefix" + + +def test_multipart_upload_roundtrip() -> None: + """Every part is uploaded and completed through independently signed requests.""" + payload = bytes(range(256)) * 40 # 10 KiB + with signed_server() as url: + store = make_store(url) + obs.put( + store, + "big", + payload, + use_multipart=True, + chunk_size=4096, + max_concurrency=4, + ) + + assert store.get("big").bytes() == payload + + assert len(select("POST", "uploads")) == 1 + assert len(select("PUT", "partNumber")) == 3, ( + "10 KiB in 4 KiB chunks is three parts" + ) + assert len(select("POST", "uploadId")) == 1 + + +@pytest.mark.asyncio +async def test_multipart_upload_async() -> None: + payload = bytes(range(256)) * 40 + with signed_server() as url: + store = make_store(url) + await obs.put_async(store, "big", payload, use_multipart=True, chunk_size=4096) + assert (await store.get_async("big")).bytes() == payload + + +def test_multipart_upload_of_empty_object() -> None: + """A completion with no buffered parts still has to produce a readable object.""" + with signed_server() as url: + store = make_store(url) + with obs.open_writer(store, "empty", buffer_size=4096) as writer: + del writer # Closed without any data written. + assert store.get("empty").bytes() == b"" + + +def test_multipart_abort_discards_upload() -> None: + """A failure mid-upload must abort, so no parts are left behind on the bucket.""" + + def failing_chunks() -> Iterator[bytes]: + yield b"x" * 8192 + msg = "input went away" + raise RuntimeError(msg) + + with signed_server() as url: + store = make_store(url) + with pytest.raises(RuntimeError, match="input went away"): + obs.put( + store, + "aborted", + failing_chunks(), + use_multipart=True, + chunk_size=4096, + ) + + assert FakeS3.uploads == {}, "the upload was aborted, not left dangling" + assert "aborted" not in FakeS3.objects + + assert len(select("DELETE", "uploadId")) == 1 + + +def test_retry_resigns_each_attempt() -> None: + """A retried request must be signed again, never reuse the previous signature.""" + with signed_server() as url: + signed: list[str] = [] + + def counting_signer(method: str, uri: str, headers: dict[str, str]): + signed.append(uri) + return signer(method, uri, headers) + + store = RemoteSignedS3Store( + url, + counting_signer, + client_options=CLIENT_OPTIONS, + retry_config=NO_WAIT_RETRIES, + ) + FakeS3.objects["chunk"] = DATA + FakeS3.fail_next = 2 + + assert store.get("chunk").bytes() == DATA + + assert len(signed) == 3, "two failed attempts plus the successful one" + + +def test_retries_exhausted_surfaces_server_error() -> None: + with signed_server() as url: + store = make_store(url, retry_config={**NO_WAIT_RETRIES, "max_retries": 1}) + FakeS3.objects["chunk"] = DATA + FakeS3.fail_next = 5 + + with pytest.raises(GenericError, match="SlowDown"): + store.get("chunk").bytes() + + +def test_attributes_and_tags_are_sent_as_signed_headers() -> None: + with signed_server() as url: + store = make_store(url) + obs.put( + store, + "chunk", + DATA, + attributes={"Content-Type": "application/json", "Content-Encoding": "gzip"}, + tags={"project": "zarr"}, + ) + + put = one(request for request in FakeS3.signed_requests if request.method == "PUT") + assert put.headers["content-type"] == "application/json" + assert put.headers["content-encoding"] == "gzip" + assert put.headers["x-amz-tagging"] == "project=zarr" + + +def test_pickle_roundtrip() -> None: + """Zarr with dask or multiprocessing pickles the store, so it must survive that.""" + with signed_server() as url: + store = make_store(url) + store.put("chunk", DATA) + + restored = pickle.loads(pickle.dumps(store)) + assert restored == store + assert restored.get("chunk").bytes() == DATA + + +def test_virtual_hosted_style_lists_from_bucket_root() -> None: + """With the bucket in the host, every path segment is part of the key prefix.""" + with signed_server() as url: + # The bucket lives in the host, so the URL carries only the key prefix. In path + # style `prefix` would have been mistaken for the bucket and dropped from keys. + origin = url.removesuffix("/bucket") + store = make_store(f"{origin}/prefix", virtual_hosted_style_request=True) + store.put("chunk", DATA) + + assert [meta["path"] for meta in store.list().collect()] == ["chunk"] + assert FakeS3.objects == {"prefix/chunk": DATA} + + +def test_async_signer_with_sync_method_errors() -> None: + """A coroutine signer cannot be awaited from the synchronous methods.""" + with signed_server() as url: + + async def async_signer(method: str, uri: str, headers: dict[str, str]): + return signer(method, uri, headers) + + store = RemoteSignedS3Store( + url, + async_signer, + client_options=CLIENT_OPTIONS, + ) + with pytest.raises(GenericError, match="signer callback failed"): + store.get("chunk").bytes() + + +def test_http_url_requires_allow_http() -> None: + with signed_server() as url: + store = RemoteSignedS3Store(url, signer) + with pytest.raises(GenericError): + store.put("chunk", DATA)