Adds AAD scope override feature.

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/AadAuthorizationTests.java

@@ -22,6 +22,8 @@
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.util.CosmosPagedFlux;
 import com.azure.identity.ClientSecretCredentialBuilder;
+import com.azure.identity.InteractiveBrowserCredential;
+import com.azure.identity.InteractiveBrowserCredentialBuilder;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -36,11 +38,10 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.lang.reflect.Field;
 import java.time.OffsetDateTime;
 import java.time.ZonedDateTime;
-import java.util.List;
-import java.util.Properties;
-import java.util.UUID;
+import java.util.*;
 
 public class AadAuthorizationTests extends TestSuiteBase {
     private final static Logger log = LoggerFactory.getLogger(AadAuthorizationTests.class);
@@ -193,6 +194,106 @@ public void createAadTokenCredential() throws InterruptedException {
         Thread.sleep(SHUTDOWN_TIMEOUT);
     }
 
+    @Test(groups = { "emulator" }, timeOut = 10 * TIMEOUT)
+    public void testAadScopeOverride_EmulatorStyleSetup() throws Exception {
+        CosmosAsyncClient setupClient = null;
+        CosmosAsyncClient aadClient = null;
+        String containerName = UUID.randomUUID().toString();
+        String overrideScope = "https://cosmos.azure.com/.default";
+
+        try {
+            setupClient = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .key(TestConfigurations.MASTER_KEY)
+                .buildAsyncClient();
+
+            setupClient.createDatabase(databaseId).block();
+            setupClient.getDatabase(databaseId).createContainer(containerName, PARTITION_KEY_PATH).block();
+        } finally {
+            if (setupClient != null) {
+                safeClose(setupClient);
+            }
+        }
+
+        Thread.sleep(TIMEOUT);
+
+        setEnv("AZURE_COSMOS_AAD_SCOPE_OVERRIDE", overrideScope);
+
+        TokenCredential emulatorCredential =
+            new AadSimpleEmulatorTokenCredential(TestConfigurations.MASTER_KEY);
+
+        aadClient = new CosmosClientBuilder()
+            .endpoint(TestConfigurations.HOST)
+            .credential(emulatorCredential)
+            .buildAsyncClient();
+
+        try {
+            CosmosAsyncContainer container = aadClient
+                .getDatabase(databaseId)
+                .getContainer(containerName);
+
+            String itemId = UUID.randomUUID().toString();
+            String pk = UUID.randomUUID().toString();
+            ItemSample item = getDocumentDefinition(itemId, pk);
+
+            container.createItem(item).block();
+
+            List<String> scopes = AadSimpleEmulatorTokenCredential.getLastScopes();
+            assert scopes != null && scopes.size() == 1;
+            assert overrideScope.equals(scopes.get(0));
+
+            container.deleteItem(item.id, new PartitionKey(item.mypk)).block();
+        } finally {
+            try {
+                CosmosAsyncClient cleanupClient = new CosmosClientBuilder()
+                    .endpoint(TestConfigurations.HOST)
+                    .key(TestConfigurations.MASTER_KEY)
+                    .buildAsyncClient();
+                try {
+                    cleanupClient.getDatabase(databaseId).delete().block();
+                } finally {
+                    safeClose(cleanupClient);
+                }
+            } finally {
+                if (aadClient != null) {
+                    safeClose(aadClient);
+                }
+                setEnv("AZURE_COSMOS_AAD_SCOPE_OVERRIDE", "");
+            }
+        }
+
+        Thread.sleep(SHUTDOWN_TIMEOUT);
+    }
+
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    private static void setEnv(String key, String value) throws Exception {
+        Map<String, String> env = System.getenv();
+        Class<?> cl = env.getClass();
+        try {
+            Field field = cl.getDeclaredField("m");
+            field.setAccessible(true);
+            Map<String, String> writableEnv = (Map<String, String>) field.get(env);
+            if (value == null) {
+                writableEnv.remove(key);
+            } else {
+                writableEnv.put(key, value);
+            }
+        } catch (NoSuchFieldException nsfe) {
+            Field[] fields = cl.getDeclaredFields();
+            for (Field f : fields) {
+                if (f.getType().getName().equals("java.util.Map")) {
+                    f.setAccessible(true);
+                    Map<String, String> map = (Map<String, String>) f.get(env);
+                    if (value == null) {
+                        map.remove(key);
+                    } else {
+                        map.put(key, value);
+                    }
+                }
+            }
+        }
+    }
+
     private ItemSample getDocumentDefinition(String itemId, String partitionKeyValue) {
         ItemSample itemSample = new ItemSample();
         itemSample.id = itemId;
@@ -224,6 +325,11 @@ class AadSimpleEmulatorTokenCredential implements TokenCredential {
         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\"]}";
 
+        private static volatile List<String> lastScopes = Collections.emptyList();
+
+        public static List<String> getLastScopes() {
+            return lastScopes;
+        }
         public AadSimpleEmulatorTokenCredential(String emulatorKey) {
             if (emulatorKey == null || emulatorKey.isEmpty()) {
                 throw new IllegalArgumentException("emulatorKey");
@@ -234,6 +340,11 @@ public AadSimpleEmulatorTokenCredential(String emulatorKey) {
 
         @Override
         public Mono<AccessToken> getToken(TokenRequestContext tokenRequestContext) {
+            List<String> scopes = tokenRequestContext.getScopes(); // List<String>, not String[]
+            lastScopes = (scopes != null && !scopes.isEmpty())
+                ? new ArrayList<>(scopes)
+                : Collections.emptyList();
+
             String aadToken = emulatorKey_based_AAD_String();
             return Mono.just(new AccessToken(aadToken, OffsetDateTime.now().plusHours(2)));
         }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java

@@ -84,6 +84,9 @@ public class Configs {
     private static final String QUERY_PLAN_RESPONSE_TIMEOUT_IN_SECONDS = "COSMOS.QUERY_PLAN_RESPONSE_TIMEOUT_IN_SECONDS";
     private static final String ADDRESS_REFRESH_RESPONSE_TIMEOUT_IN_SECONDS = "COSMOS.ADDRESS_REFRESH_RESPONSE_TIMEOUT_IN_SECONDS";
 
+    public static final String AAD_SCOPE_OVERRIDE = "AZURE_COSMOS_AAD_SCOPE_OVERRIDE";
+    private static final String DEFAULT_AAD_SCOPE_OVERRIDE = "";
+
     public static final String NON_IDEMPOTENT_WRITE_RETRY_POLICY = "COSMOS.WRITE_RETRY_POLICY";
     public static final String NON_IDEMPOTENT_WRITE_RETRY_POLICY_VARIABLE = "COSMOS_WRITE_RETRY_POLICY";
 
@@ -456,6 +459,14 @@ public int getGlobalEndpointManagerMaxInitializationTimeInSeconds() {
         return getJVMConfigAsInt(GLOBAL_ENDPOINT_MANAGER_INITIALIZATION_TIME_IN_SECONDS, DEFAULT_GLOBAL_ENDPOINT_MANAGER_INITIALIZATION_TIME_IN_SECONDS);
     }
 
+    public static String getAadScopeOverride() {
+        return System.getProperty(
+            AAD_SCOPE_OVERRIDE,
+            firstNonNull(
+                emptyToNull(System.getenv().get(AAD_SCOPE_OVERRIDE)),
+                DEFAULT_AAD_SCOPE_OVERRIDE));
+    }
+
     public URI getThinclientEndpoint() {
         String valueFromSystemProperty = System.getProperty(THINCLIENT_ENDPOINT);
         if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
@@ -1214,7 +1225,7 @@ public static String getAzureMonitorConnectionString() {
             System.getenv(APPLICATIONINSIGHTS_CONNECTION_STRING_VARIABLE)
         );
     }
-    
+
     public static EnumSet<AttributeNamingScheme> getDefaultOtelSpanAttributeNamingScheme() {
         String valueFromSystemProperty = System.getProperty(OTEL_SPAN_ATTRIBUTE_NAMING_SCHEME);
         if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -535,10 +535,11 @@ private RxDocumentClientImpl(URI serviceEndpoint,
                 hasAuthKeyResourceToken = false;
                 this.authorizationTokenProvider = null;
                 if (tokenCredential != null) {
-                    this.tokenCredentialScopes = new String[] {
-//                    AadTokenAuthorizationHelper.AAD_AUTH_TOKEN_COSMOS_WINDOWS_SCOPE,
-                        serviceEndpoint.getScheme() + "://" + serviceEndpoint.getHost() + "/.default"
-                    };
+                    String scopeOverride = configs.getAadScopeOverride();
+                    String defaultScope = serviceEndpoint.getScheme() + "://" + serviceEndpoint.getHost() + "/.default";
+                    String scopeToUse = (scopeOverride != null && !scopeOverride.isEmpty()) ? scopeOverride : defaultScope;
+
+                    this.tokenCredentialScopes = new String[] { scopeToUse };
                     this.tokenCredentialCache = new SimpleTokenCache(() -> this.tokenCredential
                         .getToken(new TokenRequestContext().addScopes(this.tokenCredentialScopes)));
                     this.authorizationTokenType = AuthorizationTokenType.AadToken;