From 7585758713272f85595e4ae6e2546ef56337e214 Mon Sep 17 00:00:00 2001
From: Benjamin <ben.amuerte@gmail.com>
Date: Fri, 3 Mar 2023 18:42:06 +0100
Subject: [PATCH] Support authentication on global region for AWS IAM

---
 .../AwsIamAuthenticationOptions.java          | 27 ++++++++++++++++--
 .../AwsIamAuthenticationUnitTests.java        | 28 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/spring-vault-core/src/main/java/org/springframework/vault/authentication/AwsIamAuthenticationOptions.java b/spring-vault-core/src/main/java/org/springframework/vault/authentication/AwsIamAuthenticationOptions.java
index 0f35084e9..b62d9e1b2 100644
--- a/spring-vault-core/src/main/java/org/springframework/vault/authentication/AwsIamAuthenticationOptions.java
+++ b/spring-vault-core/src/main/java/org/springframework/vault/authentication/AwsIamAuthenticationOptions.java
@@ -20,6 +20,7 @@
 import software.amazon.awssdk.auth.credentials.AwsCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.regions.providers.AwsRegionProvider;
 import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
 
@@ -78,8 +79,16 @@ public class AwsIamAuthenticationOptions {
 	 */
 	private final URI endpointUri;
 
+	/**
+	 * This parameter enables to sign the AWS request with the global region (us-east-1)
+	 * in case the Vault server making the proxy request is configured to use STS global
+	 * endpoint, and your application is deployed in another region.
+	 */
+	private final boolean useGlobalEndpoint;
+
 	private AwsIamAuthenticationOptions(String path, AwsCredentialsProvider credentialsProvider,
-			AwsRegionProvider regionProvider, @Nullable String role, @Nullable String serverId, URI endpointUri) {
+			AwsRegionProvider regionProvider, @Nullable String role, @Nullable String serverId, URI endpointUri,
+			boolean useGlobalEndpoint) {
 
 		this.path = path;
 		this.credentialsProvider = credentialsProvider;
@@ -87,6 +96,7 @@ private AwsIamAuthenticationOptions(String path, AwsCredentialsProvider credenti
 		this.role = role;
 		this.serverId = serverId;
 		this.endpointUri = endpointUri;
+		this.useGlobalEndpoint = useGlobalEndpoint;
 	}
 
 	/**
@@ -163,6 +173,8 @@ public static class AwsIamAuthenticationOptionsBuilder {
 		@Nullable
 		private String serverId;
 
+		private boolean useGlobalEndpoint;
+
 		private URI endpointUri = URI.create("https://sts.amazonaws.com/");
 
 		AwsIamAuthenticationOptionsBuilder() {
@@ -282,6 +294,14 @@ public AwsIamAuthenticationOptionsBuilder endpointUri(URI endpointUri) {
 			return this;
 		}
 
+		public AwsIamAuthenticationOptionsBuilder useGlobalEndpoint(Boolean useGlobalEndpoint) {
+
+			Assert.notNull(useGlobalEndpoint, "Flag useGlobalEndpoint must not be null");
+
+			this.useGlobalEndpoint = useGlobalEndpoint;
+			return this;
+		}
+
 		/**
 		 * Build a new {@link AwsIamAuthenticationOptions} instance.
 		 * @return a new {@link AwsIamAuthenticationOptions}.
@@ -289,9 +309,12 @@ public AwsIamAuthenticationOptionsBuilder endpointUri(URI endpointUri) {
 		public AwsIamAuthenticationOptions build() {
 
 			Assert.state(this.credentialsProvider != null, "Credentials or CredentialProvider must not be null");
+			if (useGlobalEndpoint) {
+				regionProvider(() -> Region.US_EAST_1);
+			}
 
 			return new AwsIamAuthenticationOptions(this.path, this.credentialsProvider, this.regionProvider, this.role,
-					this.serverId, this.endpointUri);
+					this.serverId, this.endpointUri, this.useGlobalEndpoint);
 		}
 
 	}
diff --git a/spring-vault-core/src/test/java/org/springframework/vault/authentication/AwsIamAuthenticationUnitTests.java b/spring-vault-core/src/test/java/org/springframework/vault/authentication/AwsIamAuthenticationUnitTests.java
index 17e99ed42..16dafdd45 100644
--- a/spring-vault-core/src/test/java/org/springframework/vault/authentication/AwsIamAuthenticationUnitTests.java
+++ b/spring-vault-core/src/test/java/org/springframework/vault/authentication/AwsIamAuthenticationUnitTests.java
@@ -16,8 +16,15 @@
 package org.springframework.vault.authentication;
 
 import java.time.Duration;
+import java.util.Base64;
 
+import org.hamcrest.BaseMatcher;
+import org.hamcrest.Description;
+import org.hamcrest.Matcher;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.regions.Region;
@@ -102,4 +109,25 @@ void shouldUsingAuthenticationSteps() {
 		assertThat(((LoginToken) login).isRenewable()).isTrue();
 	}
 
+	@Nested
+	@DisplayName("Unit Tests AwsIamAuthenticationOptions")
+	class AwsIamAuthenticationOptionsUnitTests {
+
+		@Test
+		void shouldSignRequestOnGlobalRegion() {
+			AwsIamAuthenticationOptions options = AwsIamAuthenticationOptions.builder().role("foo-role")
+					.regionProvider(() -> Region.US_WEST_1).credentials(AwsBasicCredentials.create("foo", "bar"))
+					.useGlobalEndpoint(true).build();
+
+			assertThat(options.getRegionProvider().getRegion()).isEqualTo(Region.US_EAST_1);
+		}
+
+		@Test
+		void shouldThrowExceptionWhenUseGlobalRegionIsNull() {
+			assertThatThrownBy(() -> AwsIamAuthenticationOptions.builder().useGlobalEndpoint(null))
+					.isInstanceOf(IllegalArgumentException.class).hasMessageContaining("useGlobalEndpoint");
+		}
+
+	}
+
 }