From 090653906e50f1a823ec135bca02ccdd23380653 Mon Sep 17 00:00:00 2001
From: Mark Rousskov <mark.simulacrum@gmail.com>
Date: Sun, 19 Nov 2023 15:55:22 -0500
Subject: [PATCH] Limit write access to specific prefix

This uses Cognito as a dispatch authority to convert OIDC claims to IAM
condition values, and then fitlers the resulting role to only writing
into the passed sha.

See https://awsteele.com/blog/2023/10/25/aws-role-session-tags-for-github-actions.html for some related context.
---
 terraform/rustc-ci/impl/artifacts.tf | 36 +++++++++++++++++++++-------
 terraform/rustc-ci/impl/cognito.tf   | 23 ++++++++++++++++++
 2 files changed, 51 insertions(+), 8 deletions(-)
 create mode 100644 terraform/rustc-ci/impl/cognito.tf

diff --git a/terraform/rustc-ci/impl/artifacts.tf b/terraform/rustc-ci/impl/artifacts.tf
index 2724cdc79..b9f9342d5 100644
--- a/terraform/rustc-ci/impl/artifacts.tf
+++ b/terraform/rustc-ci/impl/artifacts.tf
@@ -150,8 +150,12 @@ resource "aws_s3_bucket_inventory" "artifacts" {
   }
 }
 
-resource "aws_iam_role" "try_builds" {
-  name = "${var.iam_prefix}--try-role"
+data "aws_iam_openid_connect_provider" "gha" {
+  url = "https://token.actions.githubusercontent.com"
+}
+
+resource "aws_iam_role" "oidc" {
+  name = "${var.iam_prefix}--role"
 
   assume_role_policy = jsonencode({
     Version = "2012-10-17"
@@ -160,11 +164,29 @@ resource "aws_iam_role" "try_builds" {
         Effect = "Allow"
         Action = "sts:AssumeRoleWithWebIdentity"
         Principal = {
-          Federated = "arn:aws:iam::890664054962:oidc-provider/token.actions.githubusercontent.com"
+          Federated = "cognito-identity.amazonaws.com"
+        }
+        Condition = {
+          StringEquals = {
+            "cognito-identity.amazonaws.com:aud" = "${aws_cognito_identity_pool.main.id}"
+            // This forces the caller to set the session name according to the caller's run & sha
+            "sts:RoleSessionName"       = "$${aws:RequestTag/run_id}@$${aws:RequestTag/sha}"
+            "aws:RequestTag/repository" = "${var.source_repo}"
+            // For now only allow new bors & try builds
+            "aws:RequestTag/ref"        = "refs/heads/automation/bors/try"
+            "aws:RequestTag/event_name" = "push"
+          }
+        }
+      },
+      {
+        Effect = "Allow"
+        Action = "sts:TagSession"
+        Principal = {
+          Federated = "cognito-identity.amazonaws.com"
         }
         Condition = {
           StringEquals = {
-            "token.actions.githubusercontent.com:sub" = "repo:${var.repo}:ref:refs/heads/automation/bors/try"
+            "cognito-identity.amazonaws.com:aud" = "${aws_cognito_identity_pool.main.id}"
           }
         }
       }
@@ -180,10 +202,8 @@ resource "aws_iam_role" "try_builds" {
           Sid    = "ArtifactsBucketWrite"
           Effect = "Allow"
           Resource = [
-            "${aws_s3_bucket.artifacts.arn}/rustc-builds-try",
-            "${aws_s3_bucket.artifacts.arn}/rustc-builds-try/*",
-            "${aws_s3_bucket.artifacts.arn}/rustc-builds-try-alt",
-            "${aws_s3_bucket.artifacts.arn}/rustc-builds-try-alt/*",
+            "${aws_s3_bucket.artifacts.arn}/rustc-builds/$${aws:PrincipalTag/sha}/*",
+            "${aws_s3_bucket.artifacts.arn}/rustc-builds-alt/$${aws:PrincipalTag/sha}/*",
           ]
           Action = [
             "s3:GetObject",
diff --git a/terraform/rustc-ci/impl/cognito.tf b/terraform/rustc-ci/impl/cognito.tf
new file mode 100644
index 000000000..61e64d616
--- /dev/null
+++ b/terraform/rustc-ci/impl/cognito.tf
@@ -0,0 +1,23 @@
+resource "aws_cognito_identity_pool" "main" {
+  identity_pool_name               = "${var.iam_prefix}--rustc-ci"
+  allow_classic_flow               = true
+  allow_unauthenticated_identities = false
+  openid_connect_provider_arns     = ["${data.aws_iam_openid_connect_provider.gha.arn}"]
+}
+
+resource "aws_cognito_identity_pool_provider_principal_tag" "gha_mapper" {
+  identity_pool_id       = aws_cognito_identity_pool.main.id
+  identity_provider_name = data.aws_iam_openid_connect_provider.gha.arn
+  use_defaults           = false
+
+  // This maps the claims on the left (from GHA, see https://token.actions.githubusercontent.com/.well-known/openid-configuration)
+  // to "RequestTag"'s on the right. These are then matchable in the AssumeRole policy.
+  principal_tags = {
+    actor        = "actor"
+    workflow_sha = "sha"
+    run_id       = "run_id"
+    event        = "event_name"
+    ref          = "ref"
+    repository   = "repository"
+  }
+}