From 9d0468bae7e80903f50462f4c00e2831047d3c2f Mon Sep 17 00:00:00 2001
From: "justin.bertrand" <justin.bertrand@canada.ca>
Date: Mon, 9 Mar 2020 18:16:35 -0400
Subject: [PATCH] Expanding EnvironmentVaultConfiguration for
 AzureMSIAuthentication identity service configuration

Added missing configuration unit tests for AzureMSI.

Original pull request: gh-542.
---
 .../config/EnvironmentVaultConfiguration.java | 20 ++++++-
 ...rationAzureMSIAuthenticationUnitTests.java | 56 +++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)
 create mode 100644 spring-vault-core/src/test/java/org/springframework/vault/config/EnvironmentVaultConfigurationAzureMSIAuthenticationUnitTests.java

diff --git a/spring-vault-core/src/main/java/org/springframework/vault/config/EnvironmentVaultConfiguration.java b/spring-vault-core/src/main/java/org/springframework/vault/config/EnvironmentVaultConfiguration.java
index b468c819..9aed67fa 100644
--- a/spring-vault-core/src/main/java/org/springframework/vault/config/EnvironmentVaultConfiguration.java
+++ b/spring-vault-core/src/main/java/org/springframework/vault/config/EnvironmentVaultConfiguration.java
@@ -145,6 +145,10 @@ import org.springframework.web.client.RestOperations;
  * <li>Azure MSI path: {@code vault.azure-msi.azure-path} (since 2.2.1, defaults to
  * {@link AzureMsiAuthenticationOptions#DEFAULT_AZURE_AUTHENTICATION_PATH})</li>
  * <li>Role: {@code vault.azure-msi.role}</li>
+ * <li>MetadataServiceUri: {@code vault.azure-msi.metadata-service-uri} (defaults to
+ * {@link AzureMsiAuthenticationOptions#DEFAULT_INSTANCE_METADATA_SERVICE_URI})</li>
+ * <li>IdentityTokenServiceUri: {@code vault.azure-msi.identity-token-service-uri} (defaults to
+ * {@link AzureMsiAuthenticationOptions#DEFAULT_IDENTITY_TOKEN_SERVICE_URI})</li>
  * </ul>
  * <li>Client Certificate authentication
  * <ul>
@@ -174,6 +178,7 @@ import org.springframework.web.client.RestOperations;
  * @see AppIdAuthentication
  * @see AppRoleAuthentication
  * @see AwsEc2Authentication
+ * @See AzureMsiAuthentication
  * @see ClientCertificateAuthentication
  * @see CubbyholeAuthentication
  * @see KubernetesAuthentication
@@ -384,12 +389,16 @@ public class EnvironmentVaultConfiguration extends AbstractVaultConfiguration
 		String role = getProperty("vault.azure-msi.role");
 		String path = getProperty("vault.azure-msi.azure-path",
 				AzureMsiAuthenticationOptions.DEFAULT_AZURE_AUTHENTICATION_PATH);
-
+		URI metadataServiceUri = getUri("vault.azure-msi.metadata-service-uri",
+				AzureMsiAuthenticationOptions.DEFAULT_INSTANCE_METADATA_SERVICE_URI);
+		URI identityTokenServiceUri = getUri("vault.azure-msi.identity-token-service-uri",
+				AzureMsiAuthenticationOptions.DEFAULT_IDENTITY_TOKEN_SERVICE_URI);
 		Assert.hasText(role,
 				"Vault Azure MSI authentication: Role (vault.azure-msi.role) must not be empty");
 
 		AzureMsiAuthenticationOptionsBuilder builder = AzureMsiAuthenticationOptions
-				.builder().role(role).path(path);
+				.builder().role(role).path(path)
+				.instanceMetadataUri(metadataServiceUri).identityTokenServiceUri(identityTokenServiceUri);
 
 		return new AzureMsiAuthentication(builder.build(), restOperations());
 	}
@@ -435,6 +444,13 @@ public class EnvironmentVaultConfiguration extends AbstractVaultConfiguration
 		return getEnvironment().getProperty(key, defaultValue);
 	}
 
+	@Nullable
+	private URI getUri(String key, URI defaultValue) {
+
+		String value = getProperty(key);
+		return value != null ? URI.create(value) : defaultValue;
+	}
+
 	@Nullable
 	private Resource getResource(String key) {
 
diff --git a/spring-vault-core/src/test/java/org/springframework/vault/config/EnvironmentVaultConfigurationAzureMSIAuthenticationUnitTests.java b/spring-vault-core/src/test/java/org/springframework/vault/config/EnvironmentVaultConfigurationAzureMSIAuthenticationUnitTests.java
new file mode 100644
index 00000000..a8327f7d
--- /dev/null
+++ b/spring-vault-core/src/test/java/org/springframework/vault/config/EnvironmentVaultConfigurationAzureMSIAuthenticationUnitTests.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2017-2020 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.vault.config;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+import org.springframework.test.context.TestPropertySource;
+import org.springframework.test.context.junit.jupiter.SpringExtension;
+import org.springframework.vault.authentication.AwsEc2Authentication;
+import org.springframework.vault.authentication.AzureMsiAuthentication;
+import org.springframework.vault.authentication.ClientAuthentication;
+
+/**
+ * Unit tests for {@link EnvironmentVaultConfiguration} with AzureMSI authentication.
+ *
+ * @author Justin Bertrand
+ */
+@ExtendWith(SpringExtension.class)
+@TestPropertySource(properties = { "vault.uri=https://localhost:8123",
+		"vault.authentication=azure", "vault.azure-msi.role=role" })
+class EnvironmentVaultConfigurationAzureMSIAuthenticationUnitTests {
+
+	@Configuration
+	@Import(EnvironmentVaultConfiguration.class)
+	static class ApplicationConfiguration {
+	}
+
+	@Autowired
+	EnvironmentVaultConfiguration configuration;
+
+	@Test
+	void shouldConfigureAuthentication() {
+
+		ClientAuthentication clientAuthentication = configuration.clientAuthentication();
+
+		assertThat(clientAuthentication).isInstanceOf(AzureMsiAuthentication.class);
+	}
+}