CosmosDB: AAD auth related fix (#22478)

* Add the missing auth context when creating the barrier related requests.

Author: milismsft

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/BarrierRequestHelper.java

@@ -39,6 +39,10 @@ public static Mono<RxDocumentServiceRequest> createAsync(
 
         AuthorizationTokenType originalRequestTokenType = request.authorizationTokenType;
 
+        if (authorizationTokenProvider != null && authorizationTokenProvider.getAuthorizationTokenType() != null) {
+            originalRequestTokenType = authorizationTokenProvider.getAuthorizationTokenType();
+        }
+
         if (originalRequestTokenType == AuthorizationTokenType.Invalid) {
             String message = "AuthorizationTokenType not set for the read request";
             assert false : message;
@@ -53,7 +57,8 @@ public static Mono<RxDocumentServiceRequest> createAsync(
                 OperationType.HeadFeed,
                 null,
                 ResourceType.Database,
-                null);
+                null,
+                originalRequestTokenType);
         } else if (request.getIsNameBased()) {
             // Name based server request
 
@@ -63,13 +68,16 @@ public static Mono<RxDocumentServiceRequest> createAsync(
             barrierLsnRequest = RxDocumentServiceRequest.createFromName(clientContext,
                     OperationType.Head,
                     collectionLink,
-                    ResourceType.DocumentCollection);
+                    ResourceType.DocumentCollection,
+                    originalRequestTokenType);
         } else {
             // RID based Server request
             barrierLsnRequest = RxDocumentServiceRequest.create(clientContext,
                     OperationType.Head,
                     ResourceId.parse(request.getResourceId()).getDocumentCollectionId().toString(),
-                    ResourceType.DocumentCollection, null);
+                    ResourceType.DocumentCollection,
+                    null,
+                    originalRequestTokenType);
         }
 
         barrierLsnRequest.getHeaders().put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123());

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/BarrierRequestHelperTest.java

@@ -3,7 +3,12 @@
 
 package com.azure.cosmos.implementation.directconnectivity;
 
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
+import com.azure.cosmos.implementation.AuthorizationTokenType;
+import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Document;
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.IAuthorizationTokenProvider;
@@ -12,10 +17,16 @@
 import com.azure.cosmos.implementation.RxDocumentClientImpl;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.TestConfigurations;
+import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
+import reactor.core.publisher.Mono;
 
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.time.OffsetDateTime;
+import java.time.ZonedDateTime;
 import java.util.Map;
 import java.util.UUID;
 
@@ -143,6 +154,41 @@ public void barrierDocumentReadRidBasedRequest() {
         assertThat(barrierRequest.getIsNameBased()).isEqualTo(false);
     }
 
+    @Test(groups = "direct")
+    public void barrierWithAadAuthorizationTokenProviderType() throws URISyntaxException {
+
+        TokenCredential tokenCredential = new AadSimpleTokenCredential(TestConfigurations.MASTER_KEY);
+
+        IAuthorizationTokenProvider authTokenProvider = new RxDocumentClientImpl(new URI(TestConfigurations.HOST),
+                null,
+                null,
+                null,
+                null,
+                new Configs(),
+                null,
+                null,
+                tokenCredential,
+                false,
+                false,
+                false,
+                null);
+
+        ResourceType resourceType = ResourceType.DocumentCollection;
+        OperationType operationType = OperationType.Read;
+
+        Document randomResource = new Document();
+        randomResource.setId(UUID.randomUUID().toString());
+        RxDocumentServiceRequest request =
+            RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), operationType, resourceType, "/dbs/7mVFAA==/colls/7mVFAP1jpeU=", randomResource, (Map<String, String>) null);
+
+        RxDocumentServiceRequest barrierRequest = BarrierRequestHelper.createAsync(mockDiagnosticsClientContext(), request, authTokenProvider, 11l, 10l).block();
+
+        assertThat(authTokenProvider.getAuthorizationTokenType()).isEqualTo(AuthorizationTokenType.AadToken);
+        assertThat(barrierRequest.authorizationTokenType).isEqualTo(AuthorizationTokenType.AadToken);
+        assertThat(request.authorizationTokenType).isEqualTo(AuthorizationTokenType.PrimaryMasterKey);
+    }
+
+
     @DataProvider(name = "isCollectionHeadBarrierRequestArgProvider")
     public Object[][] isCollectionHeadBarrierRequestArgProvider() {
         return new Object[][]{
@@ -214,5 +260,36 @@ private Long getTargetLsn(RxDocumentServiceRequest req) {
     private Long getTargetGlobalLsn(RxDocumentServiceRequest req) {
         return Long.parseLong(getHeaderValue(req, HttpConstants.HttpHeaders.TARGET_GLOBAL_COMMITTED_LSN));
     }
+
+    class AadSimpleTokenCredential implements TokenCredential {
+        private final String keyEncoded;
+        private final String AAD_HEADER_COSMOS_EMULATOR = "{\"typ\":\"JWT\",\"alg\":\"RS256\",\"x5t\":\"CosmosEmulatorPrimaryMaster\",\"kid\":\"CosmosEmulatorPrimaryMaster\"}";
+        private final String AAD_CLAIM_COSMOS_EMULATOR_FORMAT = "{\"aud\":\"https://localhost.localhost\",\"iss\":\"https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984\",\"iat\":%d,\"nbf\":%d,\"exp\":%d,\"aio\":\"\",\"appid\":\"localhost\",\"appidacr\":\"1\",\"idp\":\"https://localhost:8081/\",\"oid\":\"96313034-4739-43cb-93cd-74193adbe5b6\",\"rh\":\"\",\"sub\":\"localhost\",\"tid\":\"EmulatorFederation\",\"uti\":\"\",\"ver\":\"1.0\",\"scp\":\"user_impersonation\",\"groups\":[\"7ce1d003-4cb3-4879-b7c5-74062a35c66e\",\"e99ff30c-c229-4c67-ab29-30a6aebc3e58\",\"5549bb62-c77b-4305-bda9-9ec66b85d9e4\",\"c44fd685-5c58-452c-aaf7-13ce75184f65\",\"be895215-eab5-43b7-9536-9ef8fe130330\"]}";
+
+        public AadSimpleTokenCredential(String emulatorKey) {
+            if (emulatorKey == null || emulatorKey.isEmpty()) {
+                throw new IllegalArgumentException("emulatorKey");
+            }
+
+            this.keyEncoded = Utils.encodeUrlBase64String(emulatorKey.getBytes());
+        }
+
+        @Override
+        public Mono<AccessToken> getToken(TokenRequestContext tokenRequestContext) {
+            String aadToken = emulatorKey_based_AAD_String();
+            return Mono.just(new AccessToken(aadToken, OffsetDateTime.now().plusHours(2)));
+        }
+
+        String emulatorKey_based_AAD_String() {
+            ZonedDateTime currentTime = ZonedDateTime.now();
+            String part1Encoded = Utils.encodeUrlBase64String(AAD_HEADER_COSMOS_EMULATOR.getBytes());
+            String part2 = String.format(AAD_CLAIM_COSMOS_EMULATOR_FORMAT,
+                currentTime.toEpochSecond(),
+                currentTime.toEpochSecond(),
+                currentTime.plusHours(2).toEpochSecond());
+            String part2Encoded = Utils.encodeUrlBase64String(part2.getBytes());
+            return part1Encoded + "." + part2Encoded + "." + this.keyEncoded;
+        }
+    }
 }