From 76a70c0a946094c4a0e49cb90a6d784ddb9370c4 Mon Sep 17 00:00:00 2001
From: samuel40791765 <sachiang@amazon.com>
Date: Wed, 24 Jul 2024 00:26:36 +0000
Subject: [PATCH] add graviton4 and specific ec2 tag termination

---
 .../codebuild/ec2_test_framework_omnibus.yaml | 27 +++++++++++++++++++
 tests/ci/cdk/cdk/components.py                | 10 ++++---
 tests/ci/run_ec2_test_framework.sh            |  3 ++-
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/tests/ci/cdk/cdk/codebuild/ec2_test_framework_omnibus.yaml b/tests/ci/cdk/cdk/codebuild/ec2_test_framework_omnibus.yaml
index 12f4cc58f3..315273f6dc 100644
--- a/tests/ci/cdk/cdk/codebuild/ec2_test_framework_omnibus.yaml
+++ b/tests/ci/cdk/cdk/codebuild/ec2_test_framework_omnibus.yaml
@@ -32,3 +32,30 @@ batch:
           EC2_INSTANCE_TYPE: "c6g.2xlarge"
           ECR_DOCKER_TAG: "amazonlinux-2023_clang-15x_sanitizer"
           TARGET_TEST_SCRIPT: "./tests/ci/run_fips_tests.sh"
+
+    # Actual tests are ran on an Graviton4 ec2 instance via SSM Commands.
+    - identifier: graviton4_tests_asan
+      buildspec: ./tests/ci/codebuild/common/run_ec2_target.yml
+      env:
+        type: LINUX_CONTAINER
+        privileged-mode: false
+        compute-type: BUILD_GENERAL1_SMALL
+        image: 620771051181.dkr.ecr.us-west-2.amazonaws.com/aws-lc-docker-images-linux-x86:ubuntu-20.04_clang-7x-bm-framework_latest
+        variables:
+          EC2_AMI: "ami-0c29a2c5cf69b5a9c"
+          EC2_INSTANCE_TYPE: "r8g.2xlarge"
+          ECR_DOCKER_TAG: "amazonlinux-2023_clang-15x_sanitizer"
+          TARGET_TEST_SCRIPT: "./tests/ci/run_posix_sanitizers.sh"
+
+    - identifier: graviton4_tests_fips
+      buildspec: ./tests/ci/codebuild/common/run_ec2_target.yml
+      env:
+        type: LINUX_CONTAINER
+        privileged-mode: false
+        compute-type: BUILD_GENERAL1_SMALL
+        image: 620771051181.dkr.ecr.us-west-2.amazonaws.com/aws-lc-docker-images-linux-x86:ubuntu-20.04_clang-7x-bm-framework_latest
+        variables:
+          EC2_AMI: "ami-0c29a2c5cf69b5a9c"
+          EC2_INSTANCE_TYPE: "r8g.2xlarge"
+          ECR_DOCKER_TAG: "amazonlinux-2023_clang-15x_sanitizer"
+          TARGET_TEST_SCRIPT: "./tests/ci/run_fips_tests.sh"
diff --git a/tests/ci/cdk/cdk/components.py b/tests/ci/cdk/cdk/components.py
index a96d13b929..8b1b9a4b57 100644
--- a/tests/ci/cdk/cdk/components.py
+++ b/tests/ci/cdk/cdk/components.py
@@ -36,8 +36,7 @@ def __init__(self, scope: Construct, id: str, *, project: codebuild.IProject, ec
                                 actions=[
                                     "codebuild:BatchGetBuildBatches",
                                     "codebuild:ListBuildBatchesForProject",
-                                    "codebuild:StopBuildBatch",
-                                    "ec2:TerminateInstances",
+                                    "codebuild:StopBuildBatch"
                                 ],
                                 resources=[project.project_arn]))
 
@@ -47,7 +46,12 @@ def __init__(self, scope: Construct, id: str, *, project: codebuild.IProject, ec
                                 actions=[
                                     "ec2:TerminateInstances",
                                 ],
-                                resources=["arn:aws:ec2:{}:{}:instance/*".format(AWS_REGION, AWS_ACCOUNT)]))
+                                resources=["arn:aws:ec2:{}:{}:instance/*".format(AWS_REGION, AWS_ACCOUNT)],
+                                conditions={
+                                    "StringEquals": {
+                                        "ec2:ResourceTag/ec2-framework-host": "ec2-framework-host"
+                                    }
+                                }))
             # ec2:Describe* API actions do not support resource-level permissions.
             lambda_function.add_to_role_policy(
                 iam.PolicyStatement(effect=iam.Effect.ALLOW,
diff --git a/tests/ci/run_ec2_test_framework.sh b/tests/ci/run_ec2_test_framework.sh
index 435e13f2f0..73def20f31 100755
--- a/tests/ci/run_ec2_test_framework.sh
+++ b/tests/ci/run_ec2_test_framework.sh
@@ -36,6 +36,7 @@ create_ec2_instances() {
     --instance-type "$2" --security-group-ids "${EC2_SECURITY_GROUP_ID}" --subnet-id "${EC2_SUBNET_ID}" \
     --block-device-mappings 'DeviceName="/dev/sda1",Ebs={DeleteOnTermination=True,VolumeSize=200}' \
     --tag-specifications 'ResourceType="instance",Tags=[{Key="Name",Value="ec2-test-'"$CODEBUILD_WEBHOOK_TRIGGER"'"},
+                                                        {Key="ec2-framework-host",Value="ec2-framework-host"},
                                                         {Key="ec-framework-commit-tag",Value="'"$CODEBUILD_SOURCE_VERSION"'"}]' \
     --iam-instance-profile Name=aws-lc-ci-ec2-test-framework-ec2-profile \
     --placement 'AvailabilityZone=us-west-2a' \
@@ -57,7 +58,7 @@ export target_test_script="$4"
 export s3_bucket_name="aws-lc-codebuild"
 
 # create the ssm documents that will be used for the various ssm commands
-ssm_prefix=$(basename "$target_test_script" .sh)
+ssm_prefix="$(echo "$ec2_instance_type" | awk -F'.' '{print $1}')_$(basename "$target_test_script" .sh)"
 generate_ssm_document_file "${ssm_prefix}"
 
 # create ec2 instances