Enforce optional required_claims on Buildkite OIDC mappings

Pipelines that run fork-PR builds can now be constrained via
required_claims in ci_providers.yml. The relay validates each
listed claim against the token before granting identity, rejecting
builds whose claims don't match (e.g. fork branches). Pipelines
without required_claims are assumed to not expose OIDC to forks.

Author: subinz1

aws/lambda/cross_repo_ci_relay/config/ci_providers.yml

@@ -3,20 +3,38 @@
 # CI platforms that don't carry a native "repository" OIDC claim need an
 # explicit mapping from their pipeline identity to a GitHub owner/repo.
 #
-# Format: one top-level key per CI provider.  Each provider section maps
-# the provider's pipeline identifier to the GitHub repo it represents.
-#
 # IMPORTANT: Buildkite mappings use immutable organization_id/pipeline_id
 # (UUIDs) rather than slugs.  Slugs are renameable and a released slug
 # can be claimed by a different organization; IDs are permanent.
 # Find IDs via: buildkite-agent oidc request-token --audience <aud>
 #
+# SECURITY: By default, any build in a mapped pipeline can authenticate.
+# If a pipeline runs fork-PR builds, use required_claims to restrict
+# which builds are authorized (e.g. only builds on specific branches).
+# Pipelines that do NOT expose OIDC to fork builds are safe without
+# required_claims.
+#
 # The Lambda fetches this file at runtime (via CI_PROVIDERS_URL) and
 # caches it in Redis, so changes here take effect without redeployment.
 
 buildkite:
-  # organization_id/pipeline_id: owner/repo  # slug (for readability)
-  # 018e4f2a-1b2c-3d4e/018e5a6b-7c8d-9e0f: vllm-project/vllm  # vllm/ci
+  # Simple format (pipeline does not run fork builds):
+  #   organization_id/pipeline_id: owner/repo
+  #
+  # Constrained format (pipeline may run fork builds):
+  #   organization_id/pipeline_id:
+  #     repo: owner/repo
+  #     required_claims:
+  #       build_branch: [main, nightly]  # only these branches can auth
+  #       cluster_id: specific-cluster-uuid
+
+  # Example entries:
+  # 018e4f2a-1b2c-3d4e/018e5a6b-7c8d-9e0f: vllm-project/vllm
+  #
+  # 018e4f2a-1b2c-3d4e/018e9a0b-1c2d-3e4f:
+  #   repo: vllm-project/vllm
+  #   required_claims:
+  #     build_branch: [main, release]
 
 # gitlab:
 #   group/project: owner/repo

aws/lambda/cross_repo_ci_relay/tests/test_jwt_helper.py

@@ -185,6 +185,57 @@ def test_buildkite_uses_correct_audience(self):
 
         self.assertEqual(self.mock_decode.call_args.kwargs["audience"], AUDIENCE)
 
+    def test_required_claims_pass_when_matching(self):
+        load_ci_provider_mappings(
+            {
+                "buildkite": {
+                    "org-uuid-123/pipe-uuid-456": {
+                        "repo": "myorg/myrepo",
+                        "required_claims": {"build_branch": ["main", "nightly"]},
+                    }
+                }
+            }
+        )
+        self.mock_decode.return_value = _fake_buildkite_claims(build_branch="main")
+        claims = verify_oidc_token("bk.oidc.token")
+        self.assertEqual(claims["repository"], "myorg/myrepo")
+
+    def test_required_claims_reject_disallowed_branch(self):
+        load_ci_provider_mappings(
+            {
+                "buildkite": {
+                    "org-uuid-123/pipe-uuid-456": {
+                        "repo": "myorg/myrepo",
+                        "required_claims": {"build_branch": ["main", "nightly"]},
+                    }
+                }
+            }
+        )
+        self.mock_decode.return_value = _fake_buildkite_claims(
+            build_branch="fork-pr-branch"
+        )
+        with self.assertRaises(HTTPException) as ctx:
+            verify_oidc_token("bk.oidc.token")
+        self.assertEqual(ctx.exception.status_code, 403)
+        self.assertIn("build_branch", ctx.exception.detail)
+
+    def test_required_claims_reject_missing_claim(self):
+        load_ci_provider_mappings(
+            {
+                "buildkite": {
+                    "org-uuid-123/pipe-uuid-456": {
+                        "repo": "myorg/myrepo",
+                        "required_claims": {"cluster_id": ["cluster-abc"]},
+                    }
+                }
+            }
+        )
+        self.mock_decode.return_value = _fake_buildkite_claims()
+        with self.assertRaises(HTTPException) as ctx:
+            verify_oidc_token("bk.oidc.token")
+        self.assertEqual(ctx.exception.status_code, 403)
+        self.assertIn("cluster_id", ctx.exception.detail)
+
 
 class TestLoadCIProviderMappings(unittest.TestCase):
     """Tests for loading CI provider repo mappings from ci_providers.yml."""
@@ -205,17 +256,40 @@ def test_loads_valid_buildkite_entries(self):
         }
         load_ci_provider_mappings(raw)
         self.assertEqual(
-            BUILDKITE_REPO_MAP[("org-id-1", "pipe-id-1")], "vllm-project/vllm"
+            BUILDKITE_REPO_MAP[("org-id-1", "pipe-id-1")]["repo"],
+            "vllm-project/vllm",
+        )
+        self.assertEqual(
+            BUILDKITE_REPO_MAP[("org-id-2", "pipe-id-2")]["repo"], "acme/repo"
         )
-        self.assertEqual(BUILDKITE_REPO_MAP[("org-id-2", "pipe-id-2")], "acme/repo")
+
+    def test_loads_constrained_entries(self):
+        raw = {
+            "buildkite": {
+                "org-id/pipe-id": {
+                    "repo": "myorg/myrepo",
+                    "required_claims": {
+                        "build_branch": ["main", "nightly"],
+                        "cluster_id": "cluster-uuid",
+                    },
+                }
+            }
+        }
+        load_ci_provider_mappings(raw)
+        entry = BUILDKITE_REPO_MAP[("org-id", "pipe-id")]
+        self.assertEqual(entry["repo"], "myorg/myrepo")
+        self.assertEqual(
+            entry["required_claims"]["build_branch"], ["main", "nightly"]
+        )
+        self.assertEqual(entry["required_claims"]["cluster_id"], ["cluster-uuid"])
 
     def test_empty_config_clears_map(self):
-        BUILDKITE_REPO_MAP[("old", "entry")] = "old/repo"
+        BUILDKITE_REPO_MAP[("old", "entry")] = {"repo": "old/repo", "required_claims": {}}
         load_ci_provider_mappings({})
         self.assertEqual(len(BUILDKITE_REPO_MAP), 0)
 
     def test_missing_buildkite_section_clears_map(self):
-        BUILDKITE_REPO_MAP[("old", "entry")] = "old/repo"
+        BUILDKITE_REPO_MAP[("old", "entry")] = {"repo": "old/repo", "required_claims": {}}
         load_ci_provider_mappings({"gitlab": {"group/proj": "org/repo"}})
         self.assertEqual(len(BUILDKITE_REPO_MAP), 0)
 
@@ -225,7 +299,9 @@ def test_skips_invalid_entries(self):
         }
         load_ci_provider_mappings(raw)
         self.assertNotIn(("noslash", ""), BUILDKITE_REPO_MAP)
-        self.assertEqual(BUILDKITE_REPO_MAP[("ok-id", "pipe-id")], "ok/repo")
+        self.assertEqual(
+            BUILDKITE_REPO_MAP[("ok-id", "pipe-id")]["repo"], "ok/repo"
+        )
 
 
 class TestUnsupportedIssuer(unittest.TestCase):

aws/lambda/cross_repo_ci_relay/utils/jwt_helper.py

@@ -35,33 +35,67 @@
 }
 
 # Runtime-populated mapping from Buildkite (organization_id, pipeline_id) to
-# the GitHub-style "owner/repo" identity.  Uses immutable IDs rather than
-# slugs to prevent identity hijacking via slug rename.  Loaded from
-# ci_providers.yml so adding a new downstream repo only requires a config
-# change — no Lambda redeployment.
-BUILDKITE_REPO_MAP: Dict[Tuple[str, str], str] = {}
+# the GitHub-style "owner/repo" identity plus optional required claims.
+# Uses immutable IDs rather than slugs to prevent identity hijacking via slug
+# rename.  Loaded from ci_providers.yml so adding a new downstream repo only
+# requires a config change — no Lambda redeployment.
+#
+# Values are dicts: {"repo": "owner/repo", "required_claims": {...}}
+# required_claims maps claim_name -> list of allowed values (any match passes).
+BUILDKITE_REPO_MAP: Dict[Tuple[str, str], dict] = {}
 
 
 def load_ci_provider_mappings(raw: dict) -> None:
     """Populate provider repo maps from parsed ci_providers.yml content.
 
-    Currently supports ``buildkite``.  Adding a new provider means adding
-    a new section reader here and a corresponding ``_extract_repo_*``
-    function below.
+    Buildkite entries support two formats:
+      org_id/pipeline_id: owner/repo              # simple
+      org_id/pipeline_id:                         # with constraints
+        repo: owner/repo
+        required_claims:
+          build_branch: [main, nightly]
     """
     BUILDKITE_REPO_MAP.clear()
     bk_section = raw.get("buildkite")
     if bk_section and isinstance(bk_section, dict):
-        for bk_key, repo in bk_section.items():
+        for bk_key, value in bk_section.items():
             bk_key_str = str(bk_key).strip()
-            repo_str = str(repo).strip()
-            if "/" not in bk_key_str or "/" not in repo_str:
-                logger.warning(
-                    "Skipping invalid buildkite entry: %s -> %s", bk_key, repo
-                )
+            if "/" not in bk_key_str:
+                logger.warning("Skipping buildkite entry without /: %s", bk_key)
                 continue
             org, pipeline = bk_key_str.split("/", 1)
-            BUILDKITE_REPO_MAP[(org, pipeline)] = repo_str
+
+            if isinstance(value, str):
+                repo_str = value.strip()
+                if "/" not in repo_str:
+                    logger.warning(
+                        "Skipping invalid buildkite entry: %s -> %s", bk_key, value
+                    )
+                    continue
+                BUILDKITE_REPO_MAP[(org, pipeline)] = {
+                    "repo": repo_str,
+                    "required_claims": {},
+                }
+            elif isinstance(value, dict):
+                repo_str = str(value.get("repo", "")).strip()
+                if "/" not in repo_str:
+                    logger.warning(
+                        "Skipping buildkite entry with invalid repo: %s", bk_key
+                    )
+                    continue
+                required = value.get("required_claims", {})
+                if not isinstance(required, dict):
+                    required = {}
+                normalized = {}
+                for k, v in required.items():
+                    if isinstance(v, list):
+                        normalized[k] = [str(x) for x in v]
+                    else:
+                        normalized[k] = [str(v)]
+                BUILDKITE_REPO_MAP[(org, pipeline)] = {
+                    "repo": repo_str,
+                    "required_claims": normalized,
+                }
     if BUILDKITE_REPO_MAP:
         logger.info(
             "Loaded %d Buildkite repo mapping(s) from ci_providers",
@@ -137,13 +171,23 @@ def _extract_repo_buildkite(claims: dict) -> str:
             401,
             "Buildkite OIDC token missing 'organization_id' or 'pipeline_id'",
         )
-    repo = BUILDKITE_REPO_MAP.get((org_id, pipeline_id))
-    if not repo:
+    entry = BUILDKITE_REPO_MAP.get((org_id, pipeline_id))
+    if not entry:
         raise HTTPException(
             403,
             f"Buildkite pipeline {org_id}/{pipeline_id} is not registered with CRCR",
         )
-    return repo
+
+    for claim_name, allowed_values in entry["required_claims"].items():
+        actual = str(claims.get(claim_name, ""))
+        if actual not in allowed_values:
+            raise HTTPException(
+                403,
+                f"Buildkite claim '{claim_name}' value '{actual}' "
+                f"not in allowed set for this pipeline",
+            )
+
+    return entry["repo"]
 
 
 _REPO_EXTRACTORS = {