Moving AuthStrippingInterceptor to registry-aws and adding region as input parameter to ContainerRegistryClient

Author: iamabhilaksh

registry/registry-aws/src/main/java/com/salesforce/multicloudj/registry/aws/AwsRegistry.java

@@ -1,6 +1,7 @@
 package com.salesforce.multicloudj.registry.aws;
 
 import com.google.auto.service.AutoService;
+import com.salesforce.multicloudj.common.aws.AwsConstants;
 import com.salesforce.multicloudj.common.aws.CommonErrorCodeMapping;
 import com.salesforce.multicloudj.common.aws.CredentialsProvider;
 import com.salesforce.multicloudj.common.exceptions.InvalidArgumentException;
@@ -9,6 +10,7 @@
 import com.salesforce.multicloudj.registry.driver.AbstractRegistry;
 import com.salesforce.multicloudj.registry.driver.OciRegistryClient;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.http.HttpRequestInterceptor;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
@@ -19,6 +21,7 @@
 
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
+import java.util.List;
 
 /**
  * AWS Elastic Container Registry (ECR) implementation.
@@ -32,7 +35,6 @@
 @AutoService(AbstractRegistry.class)
 public class AwsRegistry extends AbstractRegistry {
 
-    public static final String PROVIDER_ID = "aws";
     private static final String AWS_AUTH_USERNAME = "AWS";
 
     /** Lock for thread-safe lazy initialization of ECR client. */
@@ -41,7 +43,6 @@ public class AwsRegistry extends AbstractRegistry {
     /** Lock for thread-safe token refresh. */
     private final Object tokenLock = new Object();
 
-    private final String region;
     private final OciRegistryClient ociClient;
 
     /** Lazily initialized ECR client with double-checked locking. */
@@ -67,7 +68,6 @@ public AwsRegistry(Builder builder) {
      */
     public AwsRegistry(Builder builder, EcrClient ecrClient) {
         super(builder);
-        this.region = builder.getRegion();
         this.ecrClient = ecrClient;
         this.ociClient = registryEndpoint != null ? new OciRegistryClient(registryEndpoint, this) : null;
     }
@@ -99,6 +99,11 @@ public String getAuthToken() {
         return cachedAuthToken;
     }
 
+    @Override
+    protected List<HttpRequestInterceptor> getInterceptors() {
+        return List.of(new AuthStrippingInterceptor(registryEndpoint));
+    }
+
     /**
      * Returns true if the current time is past the halfway point of the token's validity window,
      * at which point a proactive refresh is triggered i.e. treats tokens as invalid after 50% of their lifetime.
@@ -198,18 +203,7 @@ public void close() throws Exception {
     public static final class Builder extends AbstractRegistry.Builder<AwsRegistry, Builder> {
 
         public Builder() {
-            providerId(PROVIDER_ID);
-        }
-
-        private String region;
-
-        public Builder withRegion(String region) {
-            this.region = region;
-            return this;
-        }
-
-        public String getRegion() {
-            return region;
+            providerId(AwsConstants.PROVIDER_ID);
         }
 
         @Override

registry/registry-aws/src/test/java/com/salesforce/multicloudj/registry/aws/AuthStrippingInterceptorTest.java

@@ -20,6 +20,7 @@ class AuthStrippingInterceptorTest {
     private static final String REGISTRY_ENDPOINT = "https://123456789012.dkr.ecr.us-east-1.amazonaws.com";
     private static final String REGISTRY_HOST = "123456789012.dkr.ecr.us-east-1.amazonaws.com";
     private static final String EXAMPLE_ENDPOINT_HOST = "registry.example.com";
+    private static final String S3_HOST = "s3.amazonaws.com";
     private static final String GET = "GET";
     private static final String BEARER_TOKEN = "Bearer token123";
     private static final String PATH_BLOB = "/v2/repo/blobs/sha256:abc";
@@ -78,7 +79,7 @@ void testProcess_CaseInsensitiveHostComparison() {
     @Test
     void testProcess_NoAuthHeader_NoError() {
         HttpRequest request = new BasicHttpRequest(GET, PATH_BLOB);
-        HttpContext context = contextWithHost("s3.amazonaws.com");
+        HttpContext context = contextWithHost(S3_HOST);
 
         interceptor.process(request, context);
 

registry/registry-aws/src/test/java/com/salesforce/multicloudj/registry/aws/AwsRegistryTest.java

@@ -4,6 +4,7 @@
 import com.salesforce.multicloudj.common.exceptions.SubstrateSdkException;
 import com.salesforce.multicloudj.common.exceptions.UnAuthorizedException;
 import com.salesforce.multicloudj.common.exceptions.UnknownException;
+import org.apache.http.HttpRequestInterceptor;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -19,10 +20,14 @@
 import java.time.Instant;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.List;
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.params.provider.Arguments.arguments;
 import static org.mockito.ArgumentMatchers.any;
@@ -54,7 +59,6 @@ private AwsRegistry createRegistryWithMockEcrClient(EcrClient mockEcrClient) {
         AwsRegistry.Builder builder = new AwsRegistry.Builder()
                 .withRegion(TEST_REGION)
                 .withRegistryEndpoint(TEST_REGISTRY_ENDPOINT);
-        builder.providerId(PROVIDER_ID);
         return new AwsRegistry(builder, mockEcrClient);
     }
 
@@ -81,16 +85,6 @@ private static AwsServiceException awsException(String errorCode) {
                 .build();
     }
 
-    static Stream<org.junit.jupiter.params.provider.Arguments> exceptionMappingProvider() { // NOSONAR - needed for @MethodSource resolution
-        return Stream.of(
-                arguments(new SubstrateSdkException("test"),              SubstrateSdkException.class),
-                arguments(awsException(AWS_ERROR_CODE_ACCESS_DENIED),     UnAuthorizedException.class),
-                arguments(awsException(AWS_ERROR_CODE_UNAVAILABLE),       UnknownException.class),
-                arguments(new IllegalArgumentException("invalid"),        InvalidArgumentException.class),
-                arguments(new RuntimeException("unknown"),                UnknownException.class)
-        );
-    }
-
     private void withMockedRegistry(RegistryTestAction action) throws Exception {
         EcrClient mockEcrClient = mock(EcrClient.class);
         try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
@@ -99,62 +93,78 @@ private void withMockedRegistry(RegistryTestAction action) throws Exception {
     }
 
     @Test
-    void testBuilderAndBasicProperties() throws Exception {
+    void testNoArgConstructor_CreatesInstanceWithDefaultBuilder() {
+        AwsRegistry registry = new AwsRegistry();
+        assertNotNull(registry);
+        assertNull(registry.getOciClient());
+    }
+
+    @Test
+    void testConstructor_WithBuilder_InitialisesFields() throws Exception {
         withMockedRegistry(registry -> {
-            assertNotNull(registry);
             assertEquals(PROVIDER_ID, registry.getProviderId());
             assertEquals(AUTH_USERNAME, registry.getAuthUsername());
-            assertNotNull(registry.builder());
+            assertNotNull(registry.getOciClient());
         });
     }
 
+    @Test
+    void testBuilder_InstanceMethod_ReturnsNewBuilder() throws Exception {
+        withMockedRegistry(registry -> assertNotNull(registry.builder()));
+    }
+
+    @Test
+    void testBuilder_WithRegion_GetRegion_RoundTrip() {
+        AwsRegistry.Builder builder = new AwsRegistry.Builder().withRegion(TEST_REGION);
+        assertEquals(TEST_REGION, builder.getRegion());
+    }
+
     @ParameterizedTest
     @NullAndEmptySource
     @ValueSource(strings = {"   "})
-    void testBuilder_InvalidRegion_ThrowsException(String region) {
+    void testBuilder_MissingEndpoint_ThrowsInvalidArgumentException(String endpoint) {
         assertThrows(InvalidArgumentException.class, () ->
-            new AwsRegistry.Builder()
-                    .withRegion(region)
-                    .withRegistryEndpoint(TEST_REGISTRY_ENDPOINT)
-                    .build()
-        );
+                new AwsRegistry.Builder()
+                        .withRegion(TEST_REGION)
+                        .withRegistryEndpoint(endpoint)
+                        .build());
     }
 
-    @Test
-    void testGetAuthToken_EmptyAuthorizationData_ThrowsUnknownException() throws Exception {
-        EcrClient mockEcrClient = mock(EcrClient.class);
-        GetAuthorizationTokenResponse response = GetAuthorizationTokenResponse.builder()
-                .authorizationData(Collections.emptyList())
-                .build();
-        when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class))).thenReturn(response);
+    @ParameterizedTest
+    @NullAndEmptySource
+    @ValueSource(strings = {"   "})
+    void testBuilder_MissingRegion_ThrowsInvalidArgumentException(String region) {
+        assertThrows(InvalidArgumentException.class, () ->
+                new AwsRegistry.Builder()
+                        .withRegion(region)
+                        .withRegistryEndpoint(TEST_REGISTRY_ENDPOINT)
+                        .build());
+    }
 
-        try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
-            UnknownException exception = assertThrows(UnknownException.class, registry::getAuthToken);
-            assertEquals(ERR_EMPTY_AUTH_DATA, exception.getMessage());
-        }
+    @Test
+    void testGetInterceptors_ReturnsAuthStrippingInterceptor() throws Exception {
+        withMockedRegistry(registry -> {
+            List<HttpRequestInterceptor> interceptors = registry.getInterceptors();
+            assertFalse(interceptors.isEmpty());
+            assertInstanceOf(AuthStrippingInterceptor.class, interceptors.get(0));
+        });
     }
 
     @Test
-    void testGetAuthToken_InvalidTokenFormat_ThrowsUnknownException() throws Exception {
-        EcrClient mockEcrClient = mock(EcrClient.class);
-        AuthorizationData authData = AuthorizationData.builder()
-                .authorizationToken(Base64.getEncoder().encodeToString("invalidtoken".getBytes()))
-                .expiresAt(Instant.now().plusSeconds(TOKEN_VALIDITY_SECONDS))
-                .build();
-        when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
-                .thenReturn(tokenResponse(authData));
+    void testGetOciClient_ReturnsNonNull_WhenEndpointProvided() throws Exception {
+        withMockedRegistry(registry -> assertNotNull(registry.getOciClient()));
+    }
 
-        try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
-            UnknownException exception = assertThrows(UnknownException.class, registry::getAuthToken);
-            assertEquals(ERR_INVALID_TOKEN_FORMAT, exception.getMessage());
-        }
+    @Test
+    void testGetOciClient_ReturnsNull_WhenNoEndpoint() {
+        AwsRegistry registry = new AwsRegistry();
+        assertNull(registry.getOciClient());
     }
 
     @Test
     void testGetAuthToken_TokenCachedWithinHalfwayWindow_NoRefresh() throws Exception {
         String expectedToken = "cached-token";
         EcrClient mockEcrClient = mock(EcrClient.class);
-        // Token expires 12 hours from now — halfway point is 6 hours from now, so no refresh expected
         when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
                 .thenReturn(tokenResponse(authDataWithExpiry(expectedToken, Instant.now().plusSeconds(TOKEN_VALIDITY_SECONDS))));
 
@@ -171,7 +181,6 @@ void testGetAuthToken_TokenPastHalfwayPoint_Refreshes() throws Exception {
         String refreshedToken = "refreshed-token";
         EcrClient mockEcrClient = mock(EcrClient.class);
         when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
-                // Token already past halfway point (expired 1 second ago)
                 .thenReturn(tokenResponse(authDataWithExpiry(firstToken, Instant.now().minusSeconds(1))))
                 .thenReturn(tokenResponse(authDataWithExpiry(refreshedToken, Instant.now().plusSeconds(TOKEN_VALIDITY_SECONDS))));
 
@@ -182,12 +191,39 @@ void testGetAuthToken_TokenPastHalfwayPoint_Refreshes() throws Exception {
         }
     }
 
+    @Test
+    void testGetAuthToken_EmptyAuthorizationData_ThrowsUnknownException() throws Exception {
+        EcrClient mockEcrClient = mock(EcrClient.class);
+        when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
+                .thenReturn(GetAuthorizationTokenResponse.builder().authorizationData(Collections.emptyList()).build());
+
+        try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
+            UnknownException ex = assertThrows(UnknownException.class, registry::getAuthToken);
+            assertEquals(ERR_EMPTY_AUTH_DATA, ex.getMessage());
+        }
+    }
+
+    @Test
+    void testGetAuthToken_InvalidTokenFormat_ThrowsUnknownException() throws Exception {
+        EcrClient mockEcrClient = mock(EcrClient.class);
+        AuthorizationData authData = AuthorizationData.builder()
+                .authorizationToken(Base64.getEncoder().encodeToString("invalidtoken".getBytes()))
+                .expiresAt(Instant.now().plusSeconds(TOKEN_VALIDITY_SECONDS))
+                .build();
+        when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
+                .thenReturn(tokenResponse(authData));
+
+        try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
+            UnknownException ex = assertThrows(UnknownException.class, registry::getAuthToken);
+            assertEquals(ERR_INVALID_TOKEN_FORMAT, ex.getMessage());
+        }
+    }
+
     @Test
     void testGetAuthToken_RefreshFails_FallsBackToCachedToken() throws Exception {
         String cachedToken = "still-valid-token";
         EcrClient mockEcrClient = mock(EcrClient.class);
         when(mockEcrClient.getAuthorizationToken(any(GetAuthorizationTokenRequest.class)))
-                // First call primes cache with past-halfway token, second call simulates transient failure
                 .thenReturn(tokenResponse(authDataWithExpiry(cachedToken, Instant.now().minusSeconds(1))))
                 .thenThrow(awsException(AWS_ERROR_CODE_UNAVAILABLE));
 
@@ -204,11 +240,35 @@ void testGetAuthToken_RefreshFails_NoCachedToken_ThrowsUnknownException() throws
                 .thenThrow(awsException(AWS_ERROR_CODE_UNAVAILABLE));
 
         try (AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient)) {
-            UnknownException exception = assertThrows(UnknownException.class, registry::getAuthToken);
-            assertEquals(ERR_FAILED_AUTH_TOKEN, exception.getMessage());
+            UnknownException ex = assertThrows(UnknownException.class, registry::getAuthToken);
+            assertEquals(ERR_FAILED_AUTH_TOKEN, ex.getMessage());
         }
     }
 
+    @Test
+    void testClose_WithOciAndEcrClient_ClosesAll() throws Exception {
+        EcrClient mockEcrClient = mock(EcrClient.class);
+        AwsRegistry registry = createRegistryWithMockEcrClient(mockEcrClient);
+        registry.close();
+        verify(mockEcrClient).close();
+    }
+
+    @Test
+    void testClose_WithNullEcrClient_NoError() throws Exception {
+        AwsRegistry registry = new AwsRegistry();
+        registry.close(); // should not throw
+    }
+
+    static Stream<org.junit.jupiter.params.provider.Arguments> exceptionMappingProvider() { // NOSONAR - needed for @MethodSource resolution
+        return Stream.of(
+                arguments(new SubstrateSdkException("test"),              SubstrateSdkException.class),
+                arguments(awsException(AWS_ERROR_CODE_ACCESS_DENIED),     UnAuthorizedException.class),
+                arguments(awsException(AWS_ERROR_CODE_UNAVAILABLE),       UnknownException.class),
+                arguments(new IllegalArgumentException("invalid"),        InvalidArgumentException.class),
+                arguments(new RuntimeException("unknown"),                UnknownException.class)
+        );
+    }
+
     @ParameterizedTest
     @MethodSource("exceptionMappingProvider")
     void testGetException(Throwable input, Class<? extends SubstrateSdkException> expected) throws Exception {

registry/registry-client/src/main/java/com/salesforce/multicloudj/registry/client/ContainerRegistryClient.java

@@ -70,6 +70,12 @@ public ContainerRegistryClientBuilder withRegistryEndpoint(String registryEndpoi
             return this;
         }
 
+        /** Sets the region. Required for AWS; ignored by GCP. */
+        public ContainerRegistryClientBuilder withRegion(String region) {
+            this.registryBuilder.withRegion(region);
+            return this;
+        }
+
         /** Sets a proxy endpoint override for HTTP requests. */
         public ContainerRegistryClientBuilder withProxyEndpoint(java.net.URI proxyEndpoint) {
             this.registryBuilder.withProxyEndpoint(proxyEndpoint);