Support GCP Web Identity session credentials

Signed-off-by: Connor Lafferty <[email protected]>
prometheus · Jan 31, 2024 · 28bc50e · 28bc50e
1 parent 0fc87d7
commit 28bc50e
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 1 deletion.
diff --git a/pom.xml b/pom.xml
@@ -125,6 +125,11 @@
       <version>2.2</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-iamcredentials</artifactId>
+      <version>2.0.2</version>
+    </dependency>
   </dependencies>
 
 

diff --git a/src/main/java/io/prometheus/cloudwatch/CloudWatchCollector.java b/src/main/java/io/prometheus/cloudwatch/CloudWatchCollector.java
@@ -24,6 +24,7 @@
 import org.yaml.snakeyaml.LoaderOptions;
 import org.yaml.snakeyaml.Yaml;
 import org.yaml.snakeyaml.constructor.SafeConstructor;
+import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.cloudwatch.CloudWatchClient;
@@ -39,11 +40,22 @@
 import software.amazon.awssdk.services.resourcegroupstaggingapi.model.TagFilter;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleWithWebIdentityCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest;
+
+import com.google.cloud.iam.credentials.v1.GenerateIdTokenRequest;
+import com.google.cloud.iam.credentials.v1.GenerateIdTokenResponse;
+import com.google.cloud.iam.credentials.v1.IamCredentialsClient;
 
 public class CloudWatchCollector extends Collector implements Describable {
   private static final Logger LOGGER = Logger.getLogger(CloudWatchCollector.class.getName());
 
+  private static final String SERVICE_ACCOUNT_NAME_FORMAT = "projects/-/serviceAccounts/%s";
+
+  private static final int WEB_IDENTITY_CREDENTIAL_DURATION_SECONDS = 3600;
+
+
   static class ActiveConfig {
     ArrayList<MetricRule> rules;
     CloudWatchClient cloudWatchClient;
@@ -347,12 +359,55 @@ private void loadConfig(
     }
   }
 
+  private static String getIdToken(String serviceAccountEmail) throws IOException {
+    try (IamCredentialsClient credentialsClient = IamCredentialsClient.create()) {
+      GenerateIdTokenResponse idTokenResponse =
+          credentialsClient.generateIdToken(
+              GenerateIdTokenRequest.newBuilder()
+                  .setName(String.format(SERVICE_ACCOUNT_NAME_FORMAT, serviceAccountEmail))
+                  .setAudience(serviceAccountEmail)
+                  .setIncludeEmail(true)
+                  .build());
+      return idTokenResponse.getToken();
+    }
+  }
+
   private AwsCredentialsProvider getRoleCredentialProvider(Map<String, Object> config) {
+    String roleArn = (String) config.get("role_arn");
+    if (config.containsKey("assume_role_web_identity")) {
+      // indicates we need to use gcp-based web identity to assume the role
+      return StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
+          // intentionally using anonymous credentials since this is meant to be run only in gcp
+          // environments, and the AssumeRoleWithWebIdentityRequest can run without any credentials
+          .stsClient(
+              StsClient.builder()
+                  .credentialsProvider(AnonymousCredentialsProvider.create())
+                  .region(Region.US_WEST_1)
+                  .build())
+          .refreshRequest(
+              () -> {
+                String idToken;
+                try {
+                  idToken = getIdToken((String) config.get("assume_role_web_identity"));
+                } catch (IOException e) {
+                  throw new RuntimeException(
+                      "Failed to get id token for role arn: " + roleArn, e);
+                }
+                String[] roleSplit = roleArn.split("/");
+                return AssumeRoleWithWebIdentityRequest.builder()
+                    .roleArn(roleArn)
+                    .webIdentityToken(idToken)
+                    .roleSessionName(roleSplit[roleSplit.length - 1])
+                    .durationSeconds(WEB_IDENTITY_CREDENTIAL_DURATION_SECONDS)
+                    .build();
+              })
+          .build();
+    }
     StsClient stsClient =
         StsClient.builder().region(Region.of((String) config.get("region"))).build();
     AssumeRoleRequest assumeRoleRequest =
         AssumeRoleRequest.builder()
-            .roleArn((String) config.get("role_arn"))
+            .roleArn(roleArn)
             .roleSessionName("cloudwatch_exporter")
             .build();
     return StsAssumeRoleCredentialsProvider.builder()