diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/AadAuthorizationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/AadAuthorizationTests.java index 0111623c4731..f33d1f1c72b4 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/AadAuthorizationTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/AadAuthorizationTests.java @@ -36,11 +36,15 @@ 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.ArrayList; +import java.util.Collections; +import java.util.Map; public class AadAuthorizationTests extends TestSuiteBase { private final static Logger log = LoggerFactory.getLogger(AadAuthorizationTests.class); @@ -193,6 +197,106 @@ public void createAadTokenCredential() throws InterruptedException { Thread.sleep(SHUTDOWN_TIMEOUT); } + @Test(groups = { "emulator" }, timeOut = 10 * TIMEOUT) + public void testAadScopeOverride() 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 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 env = System.getenv(); + Class cl = env.getClass(); + try { + Field field = cl.getDeclaredField("m"); + field.setAccessible(true); + Map writableEnv = (Map) 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 map = (Map) 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; @@ -219,11 +323,16 @@ public void afterMethod() { public void afterClass() { } - class AadSimpleEmulatorTokenCredential implements TokenCredential { + static class AadSimpleEmulatorTokenCredential implements TokenCredential { private final String emulatorKeyEncoded; 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 lastScopes = Collections.emptyList(); + + public static List getLastScopes() { + return lastScopes; + } public AadSimpleEmulatorTokenCredential(String emulatorKey) { if (emulatorKey == null || emulatorKey.isEmpty()) { throw new IllegalArgumentException("emulatorKey"); @@ -234,6 +343,11 @@ public AadSimpleEmulatorTokenCredential(String emulatorKey) { @Override public Mono getToken(TokenRequestContext tokenRequestContext) { + List scopes = tokenRequestContext.getScopes(); // List, 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))); } diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index efbec7b6271e..1208db3275c8 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -20,6 +20,7 @@ #### Other Changes * Added quicker cross region retry capability when a 410 `Lease Not Found` is returned by a partition in a Strong Consistency account. - See [PR 46071](https://github.com/Azure/azure-sdk-for-java/pull/46071) +* Added an option to override AAD audience scope through environment variable. See [PR 46237](https://github.com/Azure/azure-sdk-for-java/pull/46237). ### 4.73.0 (2025-07-18) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 9cbde7542afc..8a28d87f7eab 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -84,6 +84,10 @@ 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"; + private static final String AAD_SCOPE_OVERRIDE = "COSMOS.AAD_SCOPE_OVERRIDE"; + private static final String AAD_SCOPE_OVERRIDE_VARIABLE = "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"; @@ -1198,6 +1202,14 @@ public static boolean isReadAvailabilityStrategyEnabledWithPpaf() { return Boolean.parseBoolean(isReadAvailabilityStrategyEnabledWithPpaf); } + public static String getAadScopeOverride() { + return System.getProperty( + AAD_SCOPE_OVERRIDE, + firstNonNull( + emptyToNull(System.getenv().get(AAD_SCOPE_OVERRIDE_VARIABLE)), + DEFAULT_AAD_SCOPE_OVERRIDE)); + } + public static int getWarnLevelLoggingThresholdForPpaf() { String warnLevelLoggingThresholdForPpaf = System.getProperty( WARN_LEVEL_LOGGING_THRESHOLD_FOR_PPAF, @@ -1214,7 +1226,7 @@ public static String getAzureMonitorConnectionString() { System.getenv(APPLICATIONINSIGHTS_CONNECTION_STRING_VARIABLE) ); } - + public static EnumSet getDefaultOtelSpanAttributeNamingScheme() { String valueFromSystemProperty = System.getProperty(OTEL_SPAN_ATTRIBUTE_NAMING_SCHEME); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 85051d9a23f3..fa48e6dcc037 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/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;