Update secret manager retrieval (#516)

* Update secrete manager retrieval * Update log message * Run formatter
bazelbuild · Sep 10, 2020 · dcdcf3a · dcdcf3a
1 parent 5380c4a
commit dcdcf3a
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 12 deletions.
diff --git a/defs.bzl b/defs.bzl
@@ -58,6 +58,7 @@ def buildfarm_init(name="buildfarm"):
             "com.amazonaws:aws-java-sdk-secretsmanager:1.11.729",
             "com.amazonaws:aws-java-sdk-sns:1.11.729",
             "com.amazonaws:aws-java-sdk-ssm:1.11.729",
+            "com.fasterxml.jackson.core:jackson-databind:2.9.8",
             "com.github.jnr:jnr-constants:0.9.9",
             "com.github.jnr:jnr-ffi:2.1.7",
             "com.github.jnr:jffi:1.2.16",

diff --git a/src/main/java/build/buildfarm/BUILD b/src/main/java/build/buildfarm/BUILD
@@ -106,6 +106,7 @@ java_library(
         "@maven//:com_amazonaws_aws_java_sdk_core",
         "@maven//:com_amazonaws_aws_java_sdk_sns",
         "@maven//:com_amazonaws_aws_java_sdk_secretsmanager",
+        "@maven//:com_fasterxml_jackson_core_jackson_databind",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java_util",
         "@remote_apis//:build_bazel_remote_execution_v2_remote_execution_java_proto",

diff --git a/src/main/java/build/buildfarm/metrics/aws/AwsMetricsPublisher.java b/src/main/java/build/buildfarm/metrics/aws/AwsMetricsPublisher.java
@@ -30,8 +30,11 @@
 import com.amazonaws.services.sns.model.PublishRequest;
 import com.amazonaws.services.sns.model.PublishResult;
 import com.amazonaws.util.StringUtils;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.longrunning.Operation;
+import java.io.IOException;
 import java.util.Base64;
+import java.util.HashMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -40,17 +43,16 @@ public class AwsMetricsPublisher extends AbstractMetricsPublisher {
   private static AmazonSNSAsync snsClient;
 
   private String snsTopicOperations;
-  private String accessKeyId;
-  private String secretKey;
+  private String accessKeyId = null;
+  private String secretKey = null;
   private String region;
   private int snsClientMaxConnections;
 
   public AwsMetricsPublisher(MetricsConfig metricsConfig) {
     super(metricsConfig.getClusterId());
     snsTopicOperations = metricsConfig.getAwsMetricsConfig().getOperationsMetricsTopic();
     region = metricsConfig.getAwsMetricsConfig().getRegion();
-    accessKeyId = metricsConfig.getAwsMetricsConfig().getAccessKeyId();
-    secretKey = getAwsSecretKey(metricsConfig.getAwsMetricsConfig().getSecretName());
+    getAwsSecret(metricsConfig.getAwsMetricsConfig().getSecretName());
     snsClientMaxConnections = metricsConfig.getAwsMetricsConfig().getSnsClientMaxConnections();
     if (!StringUtils.isNullOrEmpty(snsTopicOperations)
         && snsClientMaxConnections > 0
@@ -114,7 +116,7 @@ public void publishMetric(String metricName, Object metricValue) {
     throw new UnsupportedOperationException();
   }
 
-  private String getAwsSecretKey(String secretName) {
+  private void getAwsSecret(String secretName) {
     AWSSecretsManager client = AWSSecretsManagerClientBuilder.standard().withRegion(region).build();
     GetSecretValueRequest getSecretValueRequest =
         new GetSecretValueRequest().withSecretId(secretName);
@@ -123,12 +125,25 @@ private String getAwsSecretKey(String secretName) {
       getSecretValueResult = client.getSecretValue(getSecretValueRequest);
     } catch (Exception e) {
       logger.severe(String.format("Could not get secret %s from AWS.", secretName));
-      return null;
+      return;
     }
+    String secret = null;
     if (getSecretValueResult.getSecretString() != null) {
-      return getSecretValueResult.getSecretString();
+      secret = getSecretValueResult.getSecretString();
     } else {
-      return new String(Base64.getDecoder().decode(getSecretValueResult.getSecretBinary()).array());
+      secret =
+          new String(Base64.getDecoder().decode(getSecretValueResult.getSecretBinary()).array());
+    }
+
+    if (secret != null) {
+      try {
+        final ObjectMapper objectMapper = new ObjectMapper();
+        final HashMap<String, String> secretMap = objectMapper.readValue(secret, HashMap.class);
+        accessKeyId = secretMap.get("access_key");
+        secretKey = secretMap.get("secret_key");
+      } catch (IOException e) {
+        logger.severe(String.format("Could not parse secret %s from AWS", secretName));
+      }
     }
   }
 }
diff --git a/src/main/protobuf/build/buildfarm/v1test/buildfarm.proto b/src/main/protobuf/build/buildfarm/v1test/buildfarm.proto
@@ -185,11 +185,9 @@ message AwsMetricsConfig {
 
   int32 sns_client_max_connections = 2;
 
-  string access_key_id = 3;
+  string secret_name = 3;
 
-  string secret_name = 4;
-
-  string region = 5;
+  string region = 4;
 }
 
 message GcpMetricsConfig {