Support the Azure VM-assigned managed identity for automatic KMS credentials (#1035)

JAVA-4706

Author: jyemin

.evergreen/.evg.yml

@@ -237,6 +237,8 @@ functions:
           AWS_TEMP_SESSION_TOKEN=$CSFLE_AWS_TEMP_SESSION_TOKEN \
           AZURE_TENANT_ID=${azure_tenant_id} AZURE_CLIENT_ID=${azure_client_id} AZURE_CLIENT_SECRET=${azure_client_secret} \
           GCP_EMAIL=${gcp_email} GCP_PRIVATE_KEY=${gcp_private_key} \
+          AZUREKMS_KEY_VAULT_ENDPOINT=${testazurekms_keyvaultendpoint} \
+          AZUREKMS_KEY_NAME=${testazurekms_keyname} \
           REQUIRE_API_VERSION=${REQUIRE_API_VERSION} \
           CRYPT_SHARED_LIB_PATH="${PROJECT_DIRECTORY}/crypt_shared/lib/mongo_crypt_v1.so" \
           .evergreen/run-tests.sh
@@ -1445,21 +1447,38 @@ tasks:
               export GCPKMS_PROJECT=${GCPKMS_PROJECT}
               export GCPKMS_ZONE=${GCPKMS_ZONE}
               export GCPKMS_INSTANCENAME=${GCPKMS_INSTANCENAME}
-              GCPKMS_CMD="MONGODB_URI='mongodb://localhost:27017' SUCCESS=true ./.evergreen/run-mongodb-fle-gcp-auto.sh" $DRIVERS_TOOLS/.evergreen/csfle/gcpkms/run-command.sh
+              GCPKMS_CMD="MONGODB_URI=mongodb://localhost:27017 PROVIDER=gcp ./.evergreen/run-fle-on-demand-credential-test.sh" $DRIVERS_TOOLS/.evergreen/csfle/gcpkms/run-command.sh
 
-    - name: "testgcpkms-fail-task"
-      # testgcpkms-fail-task runs in a non-GCE environment.
-      # It is expected to fail to obtain GCE credentials.
+    - name: testazurekms-task
       commands:
+        - command: shell.exec
+          type: setup
+          params:
+            working_dir: src
+            shell: "bash"
+            script: |
+              ${PREPARE_SHELL}
+              echo "Copying files ... begin"
+              export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+              export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+              export AZUREKMS_PRIVATEKEYPATH=/tmp/testazurekms_privatekey
+              tar czf /tmp/mongo-csharp-driver.tgz .
+              AZUREKMS_SRC=/tmp/mongo-csharp-driver.tgz AZUREKMS_DST="~/" $DRIVERS_TOOLS/.evergreen/csfle/azurekms/copy-file.sh
+              echo "Copying files ... end"
+              echo "Untarring file ... begin"
+              AZUREKMS_CMD="tar xf mongo-csharp-driver.tgz" $DRIVERS_TOOLS/.evergreen/csfle/azurekms/run-command.sh
+              echo "Untarring file ... end"
         - command: shell.exec
           type: test
           params:
             working_dir: "src"
             shell: "bash"
             script: |
               ${PREPARE_SHELL}
-              MONGODB_URI='mongodb://localhost:27017' SUCCESS=false ./.evergreen/run-mongodb-fle-gcp-auto.sh
-
+              export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+              export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+              export AZUREKMS_PRIVATEKEYPATH=/tmp/testazurekms_privatekey
+              AZUREKMS_CMD="MONGODB_URI=mongodb://localhost:27017 PROVIDER=azure AZUREKMS_KEY_VAULT_ENDPOINT=${testazurekms_keyvaultendpoint} AZUREKMS_KEY_NAME=${testazurekms_keyname} ./.evergreen/run-fle-on-demand-credential-test.sh" $DRIVERS_TOOLS/.evergreen/csfle/azurekms/run-command.sh
 axes:
   - id: version
     display_name: MongoDB Version
@@ -1694,6 +1713,47 @@ task_groups:
             $DRIVERS_TOOLS/.evergreen/csfle/gcpkms/delete-instance.sh
     tasks:
       - testgcpkms-task
+  - name: testazurekms_task_group
+    setup_group_can_fail_task: true
+    setup_group_timeout_secs: 1800 # 30 minutes
+    setup_group:
+      - func: fetch source
+      - func: prepare resources
+      - func: fix absolute paths
+      - func: make files executable
+      - command: shell.exec
+        params:
+          shell: "bash"
+          script: |
+            ${PREPARE_SHELL}
+            echo '${testazurekms_publickey}' > /tmp/testazurekms_publickey
+            echo '${testazurekms_privatekey}' > /tmp/testazurekms_privatekey
+            # Set 600 permissions on private key file. Otherwise ssh / scp may error with permissions "are too open".
+            chmod 600 /tmp/testazurekms_privatekey
+            export AZUREKMS_CLIENTID=${testazurekms_clientid}
+            export AZUREKMS_TENANTID=${testazurekms_tenantid}
+            export AZUREKMS_SECRET=${testazurekms_secret}
+            export AZUREKMS_DRIVERS_TOOLS=$DRIVERS_TOOLS
+            export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+            export AZUREKMS_PUBLICKEYPATH=/tmp/testazurekms_publickey
+            export AZUREKMS_PRIVATEKEYPATH=/tmp/testazurekms_privatekey
+            export AZUREKMS_SCOPE=${testazurekms_scope}
+            export AZUREKMS_VMNAME_PREFIX=JAVADRIVER
+            $DRIVERS_TOOLS/.evergreen/csfle/azurekms/create-and-setup-vm.sh
+      - command: expansions.update
+        params:
+          file: testazurekms-expansions.yml
+    teardown_group:
+      - command: shell.exec
+        params:
+          shell: "bash"
+          script: |
+            ${PREPARE_SHELL}
+            export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+            export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+            $DRIVERS_TOOLS/.evergreen/csfle/azurekms/delete-vm.sh
+    tasks:
+      - testazurekms-task
 
 buildvariants:
 
@@ -1932,4 +1992,11 @@ buildvariants:
   tasks:
     - name: testgcpkms_task_group
       batchtime: 20160 # Use a batchtime of 14 days as suggested by the CSFLE test README
-    - testgcpkms-fail-task
+
+- name: testazurekms-variant
+  display_name: "Azure KMS"
+  run_on:
+    - debian11-small
+  tasks:
+    - name: testazurekms_task_group
+      batchtime: 20160 # Use a batchtime of 14 days as suggested by the CSFLE test README

.evergreen/run-fle-on-demand-credential-test.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+set -o xtrace
+set -o errexit  # Exit the script with error if any of the commands fail
+
+# Supported/used environment variables:
+#       MONGODB_URI                 Set the URI, including an optional username/password to use to connect to the server
+#       PROVIDER                    Which KMS provider to test (either "gcp" or "azure")
+#       AZUREKMS_KEY_VAULT_ENDPOINT The Azure key vault endpoint for Azure integration tests
+#       AZUREKMS_KEY_NAME           The Azure key name endpoint for Azure integration tests
+
+############################################
+#            Main Program                  #
+############################################
+
+echo "Running ${PROVIDER}} Credential Acquisition Test"
+
+if ! which java ; then
+    echo "Installing java..."
+    sudo apt install openjdk-17-jdk -y
+fi
+
+./gradlew -Dorg.mongodb.test.uri="${MONGODB_URI}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.success.enabled="true" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyVaultEndpoint="${AZUREKMS_KEY_VAULT_ENDPOINT}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyName="${AZUREKMS_KEY_NAME}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.provider="${PROVIDER}" \
+ --stacktrace --debug --info  driver-sync:test --tests ClientSideEncryptionOnDemandCredentialsTest
+first=$?
+echo $first
+
+./gradlew -Dorg.mongodb.test.uri="${MONGODB_URI}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.success.enabled="true" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyVaultEndpoint="${AZUREKMS_KEY_VAULT_ENDPOINT}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyName="${AZUREKMS_KEY_NAME}" \
+ -Dorg.mongodb.test.fle.on.demand.credential.provider="${PROVIDER}" \
+ --stacktrace --debug --info  driver-reactive-streams:test --tests ClientSideEncryptionOnDemandCredentialsTest
+second=$?
+echo $second
+
+if [ $first -ne 0 ]; then
+   exit $first
+elif [ $second -ne 0 ]; then
+   exit $second
+else
+   exit 0
+fi

.evergreen/run-mongodb-fle-gcp-auto.sh

@@ -24,6 +24,8 @@ set -o errexit  # Exit the script with error if any of the commands fail
 #       AZURE_CLIENT_SECRET         The Azure client secret for client-side encryption
 #       GCP_EMAIL                   The GCP email for client-side encryption
 #       GCP_PRIVATE_KEY             The GCP private key for client-side encryption
+#       AZUREKMS_KEY_VAULT_ENDPOINT The Azure key vault endpoint for integration tests
+#       AZUREKMS_KEY_NAME           The Azure key name endpoint for integration tests
 
 AUTH=${AUTH:-noauth}
 SSL=${SSL:-nossl}
@@ -139,6 +141,9 @@ if [ "$SLOW_TESTS_ONLY" == "true" ]; then
               --stacktrace --info testSlowOnly
 else
     ./gradlew -PjavaVersion=${JAVA_VERSION} -Dorg.mongodb.test.uri=${MONGODB_URI} \
+              -Dorg.mongodb.test.fle.on.demand.credential.test.failure.enabled="true" \
+              -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyVaultEndpoint="${AZUREKMS_KEY_VAULT_ENDPOINT}" \
+              -Dorg.mongodb.test.fle.on.demand.credential.test.azure.keyName="${AZUREKMS_KEY_NAME}" \
               -Dorg.mongodb.test.awsAccessKeyId=${AWS_ACCESS_KEY_ID} -Dorg.mongodb.test.awsSecretAccessKey=${AWS_SECRET_ACCESS_KEY} \
               -Dorg.mongodb.test.tmpAwsAccessKeyId=${AWS_TEMP_ACCESS_KEY_ID} -Dorg.mongodb.test.tmpAwsSecretAccessKey=${AWS_TEMP_SECRET_ACCESS_KEY} -Dorg.mongodb.test.tmpAwsSessionToken=${AWS_TEMP_SESSION_TOKEN} \
               -Dorg.mongodb.test.azureTenantId=${AZURE_TENANT_ID} -Dorg.mongodb.test.azureClientId=${AZURE_CLIENT_ID} -Dorg.mongodb.test.azureClientSecret=${AZURE_CLIENT_SECRET} \

.evergreen/run-tests.sh

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.authentication;
+
+import com.mongodb.MongoClientException;
+import org.bson.BsonDocument;
+import org.bson.json.JsonParseException;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.mongodb.internal.authentication.HttpHelper.getHttpContents;
+
+/**
+ * Utility class for working with Azure authentication.
+ *
+ * <p>This class should not be considered a part of the public API.</p>
+ */
+public final class AzureCredentialHelper {
+    public static BsonDocument obtainFromEnvironment() {
+        String endpoint = "http://" + "169.254.169.254:80"
+                + "/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https://vault.azure.net";
+
+        Map<String, String> headers = new HashMap<>();
+        headers.put("Metadata", "true");
+        headers.put("Accept", "application/json");
+
+        String response = getHttpContents("GET", endpoint, headers);
+        try {
+            BsonDocument responseDocument = BsonDocument.parse(response);
+            if (responseDocument.containsKey("access_token")) {
+                return new BsonDocument("accessToken", responseDocument.get("access_token"));
+            } else {
+                throw new MongoClientException("The access_token is missing from Azure IMDS metadata response.");
+            }
+        } catch (JsonParseException e) {
+            throw new MongoClientException("Exception parsing JSON from Azure IMDS metadata response.", e);
+        }
+    }
+
+    private AzureCredentialHelper() {
+    }
+}

driver-core/src/main/com/mongodb/internal/authentication/AzureCredentialHelper.java

@@ -16,7 +16,7 @@
 
 package com.mongodb.internal.authentication;
 
-import com.mongodb.MongoInternalException;
+import com.mongodb.MongoClientException;
 import com.mongodb.lang.NonNull;
 
 import java.io.BufferedReader;
@@ -41,9 +41,10 @@ public static String getHttpContents(final String method, final String endpoint,
         HttpURLConnection conn = null;
         try {
             conn = (HttpURLConnection) new URL(endpoint).openConnection();
-            conn.setRequestMethod(method);
+            conn.setConnectTimeout(10000);
             conn.setReadTimeout(10000);
-            if (headers != null) {
+            conn.setRequestMethod(method);
+           if (headers != null) {
                 for (Map.Entry<String, String> kvp : headers.entrySet()) {
                     conn.setRequestProperty(kvp.getKey(), kvp.getValue());
                 }
@@ -61,7 +62,7 @@ public static String getHttpContents(final String method, final String endpoint,
                 }
             }
         } catch (IOException e) {
-            throw new MongoInternalException("Unexpected IOException", e);
+            throw new MongoClientException("Unexpected IOException from endpoint " + endpoint + ".", e);
         } finally {
             if (conn != null) {
                 conn.disconnect();

driver-core/src/main/com/mongodb/internal/authentication/HttpHelper.java

@@ -26,6 +26,7 @@
 import com.mongodb.client.model.vault.RewrapManyDataKeyOptions;
 import com.mongodb.crypt.capi.MongoCryptOptions;
 import com.mongodb.internal.authentication.AwsCredentialHelper;
+import com.mongodb.internal.authentication.AzureCredentialHelper;
 import com.mongodb.internal.authentication.GcpCredentialHelper;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
@@ -125,6 +126,10 @@ public static BsonDocument fetchCredentials(final Map<String, Map<String, Object
         if (kmsProvidersDocument.containsKey("gcp") && kmsProvidersDocument.get("gcp").asDocument().isEmpty()) {
             kmsProvidersDocument.put("gcp", GcpCredentialHelper.obtainFromEnvironment());
         }
+        if (kmsProvidersDocument.containsKey("azure") && kmsProvidersDocument.get("azure").asDocument().isEmpty()) {
+            kmsProvidersDocument.put("azure", AzureCredentialHelper.obtainFromEnvironment());
+        }
+
         return kmsProvidersDocument;
     }
 

driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java

@@ -17,12 +17,12 @@
 package com.mongodb.reactivestreams.client;
 
 import com.mongodb.ClientEncryptionSettings;
-import com.mongodb.client.AbstractClientSideEncryptionOnDemandGcpCredentialsTest;
+import com.mongodb.client.AbstractClientSideEncryptionOnDemandCredentialsTest;
 import com.mongodb.client.vault.ClientEncryption;
 import com.mongodb.reactivestreams.client.syncadapter.SyncClientEncryption;
 import com.mongodb.reactivestreams.client.vault.ClientEncryptions;
 
-public class ClientSideEncryptionOnDemandGcpCredentialsTest extends AbstractClientSideEncryptionOnDemandGcpCredentialsTest {
+public class ClientSideEncryptionOnDemandCredentialsTest extends AbstractClientSideEncryptionOnDemandCredentialsTest {
 
     @Override
     public ClientEncryption getClientEncryption(final ClientEncryptionSettings settings) {