Load azure secret from vault in JAVA code

pom.xml

@@ -150,6 +150,16 @@
             <artifactId>loki-logback-appender</artifactId>
             <version>1.2.0</version>
         </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-security-keyvault-secrets</artifactId>
+            <version>4.7.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-identity</artifactId>
+            <version>1.10.1</version>
+        </dependency>
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-s3</artifactId>

scripts/azure-cc/entrypoint.sh

@@ -2,7 +2,6 @@
 #
 # This script must be compatible with Ash (provided in eclipse-temurin Docker image) and Bash
 
-# -- set API tokens
 if [ -z "${VAULT_NAME}" ]; then
   echo "VAULT_NAME cannot be empty"
   exit 1
@@ -13,20 +12,8 @@ if [ -z "${OPERATOR_KEY_SECRET_NAME}" ]; then
   exit 1
 fi
 
-ACCESS_TOKEN=$(wget "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fvault.azure.net" -q --header "Metadata: true" -O -| jq -e -r ".access_token")
-if [ $? -ne 0 -o -z "${ACCESS_TOKEN}" ]; then
-  echo "Failed to get access token"
-  exit 1
-fi
-
-OPERATOR_KEY=$(wget "https://${VAULT_NAME}.vault.azure.net/secrets/${OPERATOR_KEY_SECRET_NAME}/?api-version=7.4" -q --header "authorization: Bearer ${ACCESS_TOKEN}" --header "content-type: application/json" -O - | jq -e -r ".value")
-if [ $? -ne 0 -o -z "${OPERATOR_KEY}" ]; then
-  echo "Failed to get operator key"
-  exit 1
-fi
-
-export core_api_token="${OPERATOR_KEY}"
-export optout_api_token="${OPERATOR_KEY}"
+export azure_vault_name="${VAULT_NAME}"
+export azure_secret_name="${OPERATOR_KEY_SECRET_NAME}"
 
 # -- locate config file
 if [ -z "${DEPLOYMENT_ENVIRONMENT}" ]; then

scripts/azure-cc/generate-deployment-artifacts.sh

@@ -67,7 +67,7 @@ if [[ $? -ne 0 ]]; then
   exit 1
 fi
 
-az confcom acipolicygen --approve-wildcards --template-file ${OUTPUT_TEMPLATE_FILE} > ${OUTPUT_POLICY_DIGEST_FILE}
+az confcom acipolicygen --approve-wildcards --debug-mode --template-file ${OUTPUT_TEMPLATE_FILE} > ${OUTPUT_POLICY_DIGEST_FILE}
 if [[ $? -ne 0 ]]; then
   echo "Failed to generate template file"
   exit 1

src/main/java/com/uid2/operator/Const.java

@@ -16,5 +16,9 @@ public class Config extends com.uid2.shared.Const.Config {
         public static final String SharingTokenExpiryProp = "sharing_token_expiry_seconds";
         public static final String EnableClientSideTokenGenerate = "client_side_token_generate";
         public static final String ValidateServiceLinks = "validate_service_links";
+
+        public static final String AzureVaultNameProp = "azure_vault_name";
+        public static final String AzureSecretNameProp = "azure_secret_name";
+
     }
 }

src/main/java/com/uid2/operator/Main.java

@@ -5,8 +5,10 @@
 import com.uid2.operator.monitoring.IStatsCollectorQueue;
 import com.uid2.operator.monitoring.OperatorMetrics;
 import com.uid2.operator.monitoring.StatsCollectorVerticle;
+import com.uid2.operator.operatorkey.OperatorKeyRetrieverFactory;
 import com.uid2.operator.service.SecureLinkValidatorService;
-import com.uid2.operator.store.*;
+import com.uid2.operator.store.CloudSyncOptOutStore;
+import com.uid2.operator.store.OptOutCloudStorage;
 import com.uid2.operator.vertx.OperatorDisableHandler;
 import com.uid2.operator.vertx.UIDOperatorVerticle;
 import com.uid2.shared.ApplicationVersion;
@@ -92,13 +94,16 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
         this.validateServiceLinks = config.getBoolean(Const.Config.ValidateServiceLinks, false);
 
         String coreAttestUrl = this.config.getString(Const.Config.CoreAttestUrlProp);
+
+        var operatorKeyRetriever = OperatorKeyRetrieverFactory.getOperatorKeyRetriever(this.config);
+        var operatorKey = operatorKeyRetriever.retrieve();
+
         DownloadCloudStorage fsStores;
         if (coreAttestUrl != null) {
-            String coreApiToken = this.config.getString(Const.Config.CoreApiTokenProp);
             Duration disableWaitTime = Duration.ofHours(this.config.getInteger(Const.Config.FailureShutdownWaitHoursProp, 120));
             this.disableHandler = new OperatorDisableHandler(disableWaitTime, Clock.systemUTC());
 
-            var clients = createUidClients(this.vertx, coreAttestUrl, coreApiToken, this.disableHandler::handleResponseStatus);
+            var clients = createUidClients(this.vertx, coreAttestUrl, operatorKey, this.disableHandler::handleResponseStatus);
             UidCoreClient coreClient = clients.getKey();
             UidOptOutClient optOutClient = clients.getValue();
             fsStores = coreClient;
@@ -134,7 +139,7 @@ public Main(Vertx vertx, JsonObject config) throws Exception {
         this.keysetProvider = new RotatingKeysetProvider(fsStores, new GlobalScope(new CloudPath(keysetMdPath)));
         String saltsMdPath = this.config.getString(Const.Config.SaltsMetadataPathProp);
         this.saltProvider = new RotatingSaltProvider(fsStores, saltsMdPath);
-        this.optOutStore = new CloudSyncOptOutStore(vertx, fsLocal, this.config);
+        this.optOutStore = new CloudSyncOptOutStore(vertx, fsLocal, this.config, operatorKey);
 
         if (this.validateServiceLinks) {
             String serviceMdPath = this.config.getString(Const.Config.ServiceMetadataPathProp);

src/main/java/com/uid2/operator/operatorkey/AzureVaultOperatorKeyRetriever.java

@@ -0,0 +1,57 @@
+package com.uid2.operator.operatorkey;
+
+import com.azure.identity.ManagedIdentityCredentialBuilder;
+import com.azure.security.keyvault.secrets.SecretClientBuilder;
+import com.google.common.base.Strings;
+import com.uid2.operator.Const;
+import io.vertx.core.json.JsonObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class AzureVaultOperatorKeyRetriever implements IOperatorKeyRetriever {
+    private static final Logger LOGGER = LoggerFactory.getLogger(AzureVaultOperatorKeyRetriever.class);
+    private final JsonObject config;
+
+    public AzureVaultOperatorKeyRetriever(JsonObject config) {
+        this.config = config;
+    }
+
+    @Override
+    public String retrieve() {
+        // Check API token field first, if it's specified, use it.
+        var tokenValue = this.config.getString(Const.Config.CoreApiTokenProp);
+
+        if (!Strings.isNullOrEmpty(tokenValue)) {
+            return tokenValue;
+        }
+
+        // Otherwise, try to load it from vault.
+        var vaultName = this.config.getString(Const.Config.AzureVaultNameProp);
+        if (Strings.isNullOrEmpty(vaultName)) {
+            throw new IllegalArgumentException(Const.Config.AzureVaultNameProp + " is null or empty");
+        }
+
+        var secretName = this.config.getString(Const.Config.AzureSecretNameProp);
+        if (Strings.isNullOrEmpty(secretName)) {
+            throw new IllegalArgumentException(Const.Config.AzureSecretNameProp + " is null or empty");
+        }
+
+        return retrieveFromAzure(vaultName, secretName);
+    }
+
+    // ManagedIdentityCredential is used here.
+    private String retrieveFromAzure(String vaultName, String secretName) {
+        String vaultUrl = "https://" + vaultName + ".vault.azure.net";
+        LOGGER.info(String.format("Load OperatorKey secret (%s) from %s", secretName, vaultUrl));
+        // Use default ExponentialBackoff retry policy
+        var secretClient = new SecretClientBuilder()
+                .vaultUrl(vaultUrl)
+                .credential(new ManagedIdentityCredentialBuilder().build())
+                .buildClient();
+
+        var retrievedSecret = secretClient.getSecret(secretName);
+
+        LOGGER.info("OperatorKey secret is loaded.");
+        return retrievedSecret.getValue();
+    }
+}

src/main/java/com/uid2/operator/operatorkey/ConfigOperatorKeyRetriever.java

@@ -0,0 +1,17 @@
+package com.uid2.operator.operatorkey;
+
+import com.uid2.operator.Const;
+import io.vertx.core.json.JsonObject;
+
+public class ConfigOperatorKeyRetriever implements IOperatorKeyRetriever {
+    private final JsonObject config;
+
+    public ConfigOperatorKeyRetriever(JsonObject config) {
+        this.config = config;
+    }
+
+    @Override
+    public String retrieve() {
+        return this.config.getString(Const.Config.CoreApiTokenProp);
+    }
+}

src/main/java/com/uid2/operator/operatorkey/IOperatorKeyRetriever.java

@@ -0,0 +1,5 @@
+package com.uid2.operator.operatorkey;
+
+public interface IOperatorKeyRetriever {
+    public String retrieve();
+}

src/main/java/com/uid2/operator/operatorkey/OperatorKeyRetrieverFactory.java

@@ -0,0 +1,15 @@
+package com.uid2.operator.operatorkey;
+
+import io.vertx.core.json.JsonObject;
+
+public class OperatorKeyRetrieverFactory {
+    public static IOperatorKeyRetriever getOperatorKeyRetriever(JsonObject config) {
+        String enclavePlatform = config.getString("enclave_platform", "");
+        switch (enclavePlatform) {
+            case "azure-cc":
+                return new AzureVaultOperatorKeyRetriever(config);
+            default:
+                return new ConfigOperatorKeyRetriever(config);
+        }
+    }
+}

src/main/java/com/uid2/operator/store/CloudSyncOptOutStore.java

@@ -19,11 +19,11 @@
 import io.vertx.core.Handler;
 import io.vertx.core.Vertx;
 import io.vertx.core.json.JsonObject;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import io.vertx.ext.web.client.HttpResponse;
 import io.vertx.ext.web.client.WebClient;
 import io.vertx.ext.web.codec.BodyCodec;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -48,7 +48,7 @@ public class CloudSyncOptOutStore implements IOptOutStore {
     private final String remoteApiPath;
     private final String remoteApiBearerToken;
 
-    public CloudSyncOptOutStore(Vertx vertx, ICloudStorage fsLocal, JsonObject jsonConfig) throws MalformedURLException {
+    public CloudSyncOptOutStore(Vertx vertx, ICloudStorage fsLocal, JsonObject jsonConfig, String operatorKey) throws MalformedURLException {
         this.fsLocal = fsLocal;
         this.webClient = WebClient.create(vertx);
 
@@ -58,7 +58,7 @@ public CloudSyncOptOutStore(Vertx vertx, ICloudStorage fsLocal, JsonObject jsonC
             this.remoteApiPort = -1 == url.getPort() ? 80 : url.getPort();
             this.remoteApiHost = url.getHost();
             this.remoteApiPath = url.getPath();
-            this.remoteApiBearerToken = "Bearer " + jsonConfig.getString(Const.Config.OptOutApiTokenProp);
+            this.remoteApiBearerToken = "Bearer " + operatorKey;
         } else {
             this.remoteApiPort = -1;
             this.remoteApiHost = null;

src/test/java/com/uid2/operator/operatorkey/AzureVaultOperatorKeyRetrieverTest.java

@@ -0,0 +1,37 @@
+package com.uid2.operator.operatorkey;
+
+import com.uid2.operator.Const;
+import io.vertx.core.json.JsonObject;
+import org.junit.Assert;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.Assert.assertEquals;
+
+class AzureVaultOperatorKeyRetrieverTest {
+    @Test
+    public void testReturnApiTokenIfSpecified() {
+        var OPERATOR_KEY = "operator_key";
+        var config = new JsonObject().put(Const.Config.CoreApiTokenProp, OPERATOR_KEY);
+
+        var sut = new AzureVaultOperatorKeyRetriever(config);
+        var key = sut.retrieve();
+
+        assertEquals(OPERATOR_KEY, key);
+    }
+
+    @Test
+    public void testArgumentCheck_NoVaultName() {
+        var config = new JsonObject();
+
+        var sut = new AzureVaultOperatorKeyRetriever(config);
+        Assert.assertThrows(IllegalArgumentException.class, () -> sut.retrieve());
+    }
+
+    @Test
+    public void testArgumentCheck_NoSecretName() {
+        var config = new JsonObject().put(Const.Config.AzureVaultNameProp, "dummy");
+
+        var sut = new AzureVaultOperatorKeyRetriever(config);
+        Assert.assertThrows(IllegalArgumentException.class, () -> sut.retrieve());
+    }
+}