from __future__ import annotations import datetime as dt import hashlib import hmac import time import urllib.error import urllib.parse import urllib.request from dataclasses import dataclass from pathlib import Path from unittest.mock import patch class R2Error(RuntimeError): pass class R2PreconditionFailed(R2Error): pass @dataclass(frozen=True) class R2Credentials: access_key_id: str secret_access_key: str session_token: str | None = None def __post_init__(self) -> None: if not self.access_key_id or not self.secret_access_key: raise R2Error("R2 access key ID and secret access key are required") def _sign(key: bytes, message: str) -> bytes: return hmac.new(key, message.encode(), hashlib.sha256).digest() def _signing_key(secret: str, date: str, region: str) -> bytes: date_key = _sign(f"AWS4{secret}".encode(), date) region_key = _sign(date_key, region) service_key = _sign(region_key, "s3") return _sign(service_key, "aws4_request") class R2Client: """Minimal Cloudflare R2 S3 client for immutable object publication.""" def __init__( self, *, endpoint: str, bucket: str, credentials: R2Credentials, region: str = "auto", retry_delay: float = 1.0, sleep=time.sleep, timeout: float = 15.0, ) -> None: parsed = urllib.parse.urlparse(endpoint.rstrip("/")) if parsed.scheme != "https" or not parsed.netloc or parsed.path not in {"", "/"}: raise R2Error("R2 endpoint must be an HTTPS origin without a path") if parsed.username or parsed.password or parsed.query or parsed.fragment: raise R2Error("R2 endpoint must not contain credentials, query, or fragment") if not bucket or "/" in bucket: raise R2Error("R2 bucket must be a non-empty bucket name") self.endpoint = endpoint.rstrip("/") self.bucket = bucket self.credentials = credentials self.region = region self.retry_delay = retry_delay self.sleep = sleep self.timeout = timeout def put_file( self, *, key: str, file: Path, content_type: str = "application/json", cache_control: str = "public, max-age=31536000, immutable", if_none_match: bool = True, ) -> None: if not file.is_file() or file.stat().st_size == 0: raise R2Error(f"R2 upload source must be a non-empty file: {file}") self.put_bytes( key=key, body=file.read_bytes(), content_type=content_type, cache_control=cache_control, if_none_match=if_none_match, ) def put_bytes( self, *, key: str, body: bytes, content_type: str = "application/json", cache_control: str = "public, max-age=31536000, immutable", if_none_match: bool = True, ) -> None: if not key or key.startswith("/") or ".." in key.split("/"): raise R2Error("R2 object key must be a safe relative key") if not body: raise R2Error("R2 object body must not be empty") now = dt.datetime.now(dt.timezone.utc) amz_date = now.strftime("%Y%m%dT%H%M%SZ") date = now.strftime("%Y%m%d") payload_hash = hashlib.sha256(body).hexdigest() canonical_uri = "/" + urllib.parse.quote(f"{self.bucket}/{key}", safe="/-_.~") url = f"{self.endpoint}{canonical_uri}" host = urllib.parse.urlparse(self.endpoint).netloc headers = { "cache-control": cache_control, "content-type": content_type, "host": host, "x-amz-content-sha256": payload_hash, "x-amz-date": amz_date, } if if_none_match: headers["if-none-match"] = "*" if self.credentials.session_token: headers["x-amz-security-token"] = self.credentials.session_token signed_headers = ";".join(sorted(headers)) canonical_headers = "".join(f"{name}:{headers[name].strip()}\n" for name in sorted(headers)) canonical_request = "\n".join( ["PUT", canonical_uri, "", canonical_headers, signed_headers, payload_hash] ) scope = f"{date}/{self.region}/s3/aws4_request" string_to_sign = "\n".join( [ "AWS4-HMAC-SHA256", amz_date, scope, hashlib.sha256(canonical_request.encode()).hexdigest(), ] ) signature = hmac.new( _signing_key(self.credentials.secret_access_key, date, self.region), string_to_sign.encode(), hashlib.sha256, ).hexdigest() headers["authorization"] = ( f"AWS4-HMAC-SHA256 Credential={self.credentials.access_key_id}/{scope}, " f"SignedHeaders={signed_headers}, Signature={signature}" ) request = urllib.request.Request(url, data=body, headers=headers, method="PUT") for attempt in range(2): try: with urllib.request.urlopen(request, timeout=self.timeout) as response: if response.status not in {200, 201}: raise R2Error(f"R2 PUT returned HTTP {response.status}") return except urllib.error.HTTPError as error: if error.code == 412: raise R2PreconditionFailed(f"R2 object already exists: {key}") from error if attempt == 0 and error.code in {408, 429, 500, 502, 503, 504}: self.sleep(self.retry_delay) continue detail = error.read(2048).decode("utf-8", "replace") raise R2Error(f"R2 PUT failed with HTTP {error.code}: {detail}") from error except (urllib.error.URLError, TimeoutError) as error: if attempt == 0: self.sleep(self.retry_delay) continue reason = error.reason if isinstance(error, urllib.error.URLError) else str(error) raise R2Error(f"R2 PUT failed: {reason}") from error def self_check() -> None: captured: list[urllib.request.Request] = [] class Response: status = 200 def __enter__(self): return self def __exit__(self, *_args) -> None: return None def open_request(request: urllib.request.Request, **_kwargs): captured.append(request) return Response() client = R2Client( endpoint="https://account.r2.cloudflarestorage.com", bucket="results", credentials=R2Credentials("access", "secret"), ) with patch("urllib.request.urlopen", open_request): client.put_bytes(key="path/result.json", body=b"{}") if len(captured) != 1: raise R2Error("R2 self-check did not issue exactly one request") request = captured[0] if request.full_url != "https://account.r2.cloudflarestorage.com/results/path/result.json": raise R2Error("R2 self-check produced an unexpected URL") if request.get_header("If-none-match") != "*": raise R2Error("R2 self-check omitted the immutable write precondition") authorization = request.get_header("Authorization") or "" if not authorization.startswith("AWS4-HMAC-SHA256 Credential=access/"): raise R2Error("R2 self-check produced an invalid authorization header") attempts = 0 def timeout_once(request: urllib.request.Request, **_kwargs): nonlocal attempts attempts += 1 if attempts == 1: raise urllib.error.URLError(TimeoutError("timed out")) return Response() retrying = R2Client( endpoint="https://account.r2.cloudflarestorage.com", bucket="results", credentials=R2Credentials("access", "secret"), retry_delay=0, sleep=lambda _seconds: None, ) with patch("urllib.request.urlopen", timeout_once): retrying.put_bytes(key="path/retry.json", body=b"{}") if attempts != 2: raise R2Error("R2 self-check did not retry one transient write failure")