BearerTokenChallengeAuthorizationPolicy extracting tenant ID

sdk/storage/azure-storage-common/pom.xml

@@ -93,6 +93,18 @@
       <version>1.15.2</version> <!-- {x-version-update;com.azure:azure-identity;dependency} -->
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-inline</artifactId>
+      <version>4.11.0</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageBearerTokenChallengeAuthorizationPolicy.java

@@ -12,6 +12,8 @@
 import com.azure.core.util.CoreUtils;
 import reactor.core.publisher.Mono;
 
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -65,15 +67,34 @@ public Mono<Boolean> authorizeRequestOnChallenge(HttpPipelineCallContext context
         Map<String, String> challenges = extractChallengeAttributes(authHeader, BEARER_TOKEN_PREFIX);
 
         String scope = challenges.get("resource_id");
-        if (scope != null) {
+        String authorization = challenges.get("authorization_uri");
+
+        if (scope != null && authorization != null) {
             scope += DEFAULT_SCOPE;
             scopes = new String[] { scope };
             scopes = getScopes(context, scopes);
-            return setAuthorizationHeader(context, new TokenRequestContext().addScopes(scopes)).thenReturn(true);
+
+            String tenantId = extractTenantIdFromUri(authorization);
+            TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopes).setTenantId(tenantId);
+
+            return setAuthorizationHeader(context, tokenRequestContext).thenReturn(true);
         }
         return Mono.just(false);
     }
 
+    private String extractTenantIdFromUri(String uri) {
+        try {
+            String[] segments = new URI(uri).getPath().split("/");
+            if (segments.length > 1) {
+                return segments[1];
+            } else {
+                throw new RuntimeException("Invalid authorization URI: tenantId not found");
+            }
+        } catch (URISyntaxException e) {
+            throw new RuntimeException("Invalid authorization URI", e);
+        }
+    }
+
     @Override
     public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) {
         String authHeader = response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE);

sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageBearerTokenChallengeAuthorizationPolicyTests.java

@@ -0,0 +1,116 @@
+package com.azure.storage.common.policy;
+
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import static org.mockito.Mockito.*;
+
+public class StorageBearerTokenChallengeAuthorizationPolicyTests {
+
+    private String[] scopes;
+    private TokenCredential mockCredential;
+
+    @BeforeEach
+    public void setup() {
+        scopes = new String[]{"https://storage.azure.com/.default"};
+        mockCredential = mock(TokenCredential.class);
+    }
+
+    @Test
+    public void usesTokenProvidedByCredentials() {
+        StorageBearerTokenChallengeAuthorizationPolicy policy = new StorageBearerTokenChallengeAuthorizationPolicy(mockCredential, scopes);
+
+        HttpPipelineCallContext mockContext = mock(HttpPipelineCallContext.class);
+        HttpResponse mockResponse = mock(HttpResponse.class);
+        when(mockResponse.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE)).thenReturn(null);
+
+        Mono<Boolean> result = policy.authorizeRequestOnChallenge(mockContext, mockResponse);
+
+        StepVerifier.create(result)
+            .expectNext(false)
+            .verifyComplete();
+    }
+
+    @Test
+    public void doesNotSendUnauthorizedRequestWhenEnableTenantDiscoveryIsFalse() {
+        StorageBearerTokenChallengeAuthorizationPolicy policy = new StorageBearerTokenChallengeAuthorizationPolicy(mockCredential, scopes);
+
+        HttpPipelineCallContext mockContext = mock(HttpPipelineCallContext.class);
+        HttpResponse mockResponse = mock(HttpResponse.class);
+        when(mockResponse.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE)).thenReturn(null);
+
+        for (int i = 0; i < 10; i++) {
+            Mono<Boolean> result = policy.authorizeRequestOnChallenge(mockContext, mockResponse);
+            StepVerifier.create(result)
+                .expectNext(false)
+                .verifyComplete();
+        }
+    }
+
+    @Test
+    public void sendsUnauthorizedRequestWhenEnableTenantDiscoveryIsTrue() {
+        StorageBearerTokenChallengeAuthorizationPolicy realPolicy =
+            new StorageBearerTokenChallengeAuthorizationPolicy(mockCredential, scopes);
+
+        // Spy on the real instance
+        StorageBearerTokenChallengeAuthorizationPolicy policy = spy(realPolicy);
+
+        String expectedTenantId = "72f988bf-86f1-41af-91ab-2d7cd011db47";
+
+        HttpRequest httpRequest = new HttpRequest(HttpMethod.GET, "https://example.com");
+        HttpPipelineCallContext mockContext = mock(HttpPipelineCallContext.class);
+        when(mockContext.getHttpRequest()).thenReturn(httpRequest);
+
+        HttpResponse mockResponse = mock(HttpResponse.class);
+        when(mockResponse.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE)).thenReturn(
+            "Bearer authorization_uri=https://login.microsoftonline.com/" + expectedTenantId + "/oauth2/authorize resource_id=https://storage.azure.com");
+
+        // Properly stub the method on the spy
+        doReturn(Mono.empty()).when(policy).setAuthorizationHeader(any(), any());
+
+        for (int i = 0; i < 10; i++) {
+            Mono<Boolean> result = policy.authorizeRequestOnChallenge(mockContext, mockResponse);
+
+            StepVerifier.create(result)
+                .expectNext(true)  // Expect the Mono<Boolean> to emit 'true'
+                .verifyComplete();
+        }
+    }
+
+
+    @Test
+    public void usesScopeFromBearerChallenge() {
+        StorageBearerTokenChallengeAuthorizationPolicy realPolicy =
+            new StorageBearerTokenChallengeAuthorizationPolicy(mockCredential, "https://disk.compute.azure.com/.default");
+
+        // Spy on the real instance to allow stubbing setAuthorizationHeader
+        StorageBearerTokenChallengeAuthorizationPolicy policy = spy(realPolicy);
+
+        String serviceChallengeResponseScope = "https://storage.azure.com";
+
+        HttpPipelineCallContext mockContext = mock(HttpPipelineCallContext.class);
+        HttpResponse mockResponse = mock(HttpResponse.class);
+        when(mockResponse.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE)).thenReturn(
+            "Bearer authorization_uri=https://login.microsoftonline.com/72f988bf-86f1-41af-91ab-2d7cd011db47/oauth2/authorize resource_id=" + serviceChallengeResponseScope);
+
+        // Stub the setAuthorizationHeader method so it returns a completed Mono
+        doReturn(Mono.empty()).when(policy).setAuthorizationHeader(any(), any());
+
+        for (int i = 0; i < 2; i++) {
+            Mono<Boolean> result = policy.authorizeRequestOnChallenge(mockContext, mockResponse);
+
+            StepVerifier.create(result)
+                .expectNext(true)  // Expect 'true' instead of failing
+                .verifyComplete();
+        }
+    }
+
+}