ci: use Buildkite OIDC for macOS tarball uploads

ci/deploy_mz-debug/pipeline.template.yml

@@ -34,6 +34,8 @@ steps:
     command: bin/pyactivate -m ci.deploy_mz-debug.macos
     agents:
       queue: mac-aarch64
+    env:
+      MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy
     timeout_in_minutes: 30
     concurrency: 1
     concurrency_group: deploy-mz-debug/macos/aarch64

ci/deploy_mz-deploy/pipeline.template.yml

@@ -34,6 +34,8 @@ steps:
     command: bin/pyactivate -m ci.deploy_mz-deploy.macos
     agents:
       queue: mac-aarch64
+    env:
+      MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy
     timeout_in_minutes: 30
     concurrency: 1
     concurrency_group: deploy-mz-deploy/macos/aarch64

ci/deploy_mz/pipeline.template.yml

@@ -42,9 +42,12 @@ steps:
     concurrency: 1
     concurrency_group: deploy-mz/docker
 
-  - command: bin/pyactivate -m ci.deploy_mz.macos
+  - id: deploy_mz-macos-aarch64
+    command: bin/pyactivate -m ci.deploy_mz.macos
     agents:
       queue: mac-aarch64
+    env:
+      MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy
     timeout_in_minutes: 30
     concurrency: 1
     concurrency_group: deploy-mz/macos/aarch64

ci/tarball_uploader.py

@@ -10,10 +10,12 @@
 # This file contains a class that can be used to upload a tarball to the materialize-binaries S3 bucket.
 
 import os
+import subprocess
 import tarfile
 import tempfile
 import time
 from pathlib import Path
+from typing import Any
 
 import boto3
 import humanize
@@ -23,12 +25,55 @@
 
 BINARIES_BUCKET = "materialize-binaries"
 
+# When set, uploads use short-lived credentials for this role, obtained by
+# exchanging the job's Buildkite OIDC token, rather than the ambient AWS
+# credentials.
+UPLOAD_ROLE_ARN_ENV = "MZ_BINARIES_UPLOAD_ROLE_ARN"
+
+
+def _s3_client() -> Any:
+    role_arn = os.environ.get(UPLOAD_ROLE_ARN_ENV)
+    if not role_arn:
+        return boto3.client("s3")
+
+    # Requested immediately before use: OIDC tokens are short-lived and the
+    # build that precedes the upload can outlast one.
+    token = subprocess.check_output(
+        [
+            "buildkite-agent",
+            "oidc",
+            "request-token",
+            "--audience",
+            "sts.amazonaws.com",
+        ],
+        text=True,
+    ).strip()
+    session = boto3.Session()
+    creds = session.client("sts").assume_role_with_web_identity(
+        RoleArn=role_arn,
+        RoleSessionName=f"buildkite-{os.environ['BUILDKITE_JOB_ID']}"[:64],
+        WebIdentityToken=token,
+        DurationSeconds=3600,
+    )["Credentials"]
+    return boto3.client(
+        "s3",
+        aws_access_key_id=creds["AccessKeyId"],
+        aws_secret_access_key=creds["SecretAccessKey"],
+        aws_session_token=creds["SessionToken"],
+    )
+
 
 class TarballUploader:
     def __init__(self, package_name: str, version: TypedVersionBase):
         self.package_name = package_name
         self.version = version
         self.version_str = f"v{version.str_without_prefix()}"
+        self._s3: Any = None
+
+    def _client(self) -> Any:
+        if self._s3 is None:
+            self._s3 = _s3_client()
+        return self._s3
 
     def _create_tardir(self, name: str) -> tarfile.TarInfo:
         """Takes a dir path and creates a TarInfo. Sets the
@@ -46,7 +91,7 @@ def _sanitize_tarinfo(self, tarinfo: tarfile.TarInfo) -> tarfile.TarInfo:
 
     def _upload_redirect(self, key: str, to: str) -> None:
         with tempfile.NamedTemporaryFile() as empty:
-            boto3.client("s3").upload_fileobj(
+            self._client().upload_fileobj(
                 Fileobj=empty,
                 Bucket=BINARIES_BUCKET,
                 Key=key,
@@ -56,7 +101,7 @@ def _upload_redirect(self, key: str, to: str) -> None:
     def _upload_tarball(self, tarball: Path, platform: str) -> None:
         """Uploads the tarball (.tar) to the S3 bucket."""
         s3_object = f"{self.package_name}-{self.version_str}-{platform}.tar.gz"
-        boto3.client("s3").upload_file(
+        self._client().upload_file(
             Filename=str(tarball),
             Bucket=BINARIES_BUCKET,
             Key=s3_object,