Allow override aad auth endpoint of cosmos db's auth config (Azure#45016)

* Override aad auth endpoint cosmos db kafka connector

* Update unit test

* Update information of config key

* Refactor & add unit tests for config

* Fix

* Revert the wrong change

* Update changelog & doc

* Remove azure.cosmos.throughputControl.auth.aad.authEndpointOverride

* Remove

* Remove unused method & update doc

* Update sdk/cosmos/azure-cosmos-kafka-connect/CHANGELOG.md

---------

Co-authored-by: Fabian Meiswinkel <[email protected]>

Author: chanhanyi0923

sdk/cosmos/azure-cosmos-kafka-connect/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added `authEndpointOverride` option for all AAD authentication types - See [PR 45016](https://github.com/Azure/azure-sdk-for-java/pull/45016)
 
 ### 2.2.0 (2025-02-20)
 

sdk/cosmos/azure-cosmos-kafka-connect/docs/configuration-reference.md

@@ -11,6 +11,7 @@
 | `azure.cosmos.account.key`                                        | `""`             | Cosmos DB Account Key (only required in case of `auth.type` as `MasterKey`)                                                                                                                                                                                                                                                                                                                                                                                                                       |
 | `azure.cosmos.auth.aad.clientId`                                  | `""`             | The clientId/ApplicationId of the service principal. Required for `ServicePrincipal` authentication.                                                                                                                                                                                                                                                                                                                                                                                              |
 | `azure.cosmos.auth.aad.clientSecret`                              | `""`             | The client secret/password of the service principal.                                                                                                                                                                                                                                                                                                                                                                                                                                              |
+| `azure.cosmos.auth.aad.authEndpointOverride`                      | `""`             | Overrides the Azure Active Directory (AAD) authentication endpoint. This is useful when the Cosmos DB account resides in a non-public Azure cloud environment such as Azure China, Azure US Government, or Azure Germany. By default, the SDK uses the standard AAD endpoint for the public Azure cloud. Set this value if your deployment requires a custom authority URI (e.g., https://login.chinacloudapi.cn/).                                                                               |
 | `azure.cosmos.mode.gateway`                                       | `false`          | Flag to indicate whether to use gateway mode. By default it is false, means SDK uses direct mode. https://learn.microsoft.com/azure/cosmos-db/nosql/sdk-connection-modes                                                                                                                                                                                                                                                                                                                          |
 | `azure.cosmos.preferredRegionList`                                | `[]`             | Preferred regions list to be used for a multi region Cosmos DB account. This is a comma separated value (e.g., `[East US, West US]` or `East US, West US`) provided preferred regions will be used as hint. You should use a collocated kafka cluster with your Cosmos DB account and pass the kafka cluster region as preferred region. See list of azure regions [here](https://docs.microsoft.com/dotnet/api/microsoft.azure.documents.locationnames?view=azure-dotnet&preserve-view=true).    |
 | `azure.cosmos.application.name`                                   | `""`             | Application name. Will be added as the userAgent suffix.                                                                                                                                                                                                                                                                                                                                                                                                                                          |

sdk/cosmos/azure-cosmos-kafka-connect/src/main/java/com/azure/cosmos/kafka/connect/implementation/CosmosAadAuthConfig.java

@@ -5,21 +5,36 @@
 
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
 
 public class CosmosAadAuthConfig implements CosmosAuthConfig {
+    private static final Map<CosmosAzureEnvironment, String> ACTIVE_DIRECTORY_ENDPOINT_MAP;
+    static {
+        // for now we maintain a static list within the SDK these values do not change very frequently
+        ACTIVE_DIRECTORY_ENDPOINT_MAP = new HashMap<>();
+        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE, "https://login.microsoftonline.com/");
+        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_CHINA, "https://login.chinacloudapi.cn/");
+        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_US_GOVERNMENT, "https://login.microsoftonline.us/");
+        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_GERMANY, "https://login.microsoftonline.de/");
+    }
+
     private final String clientId;
     private final String clientSecret;
+    private final String authEndpointOverride;
     private final String tenantId;
     private final CosmosAzureEnvironment azureEnvironment;
 
-    public CosmosAadAuthConfig(String clientId, String clientSecret, String tenantId, CosmosAzureEnvironment azureEnvironment) {
+    public CosmosAadAuthConfig(String clientId, String clientSecret, String authEndpointOverride, String tenantId, CosmosAzureEnvironment azureEnvironment) {
         checkArgument(StringUtils.isNotEmpty(clientId), "Argument 'clientId' should not be null");
         checkArgument(StringUtils.isNotEmpty(clientSecret), "Argument 'clientSecret' should not be null");
         checkArgument(StringUtils.isNotEmpty(tenantId), "Argument 'tenantId' should not be null");
 
         this.clientId = clientId;
         this.clientSecret = clientSecret;
+        this.authEndpointOverride = authEndpointOverride;
         this.tenantId = tenantId;
         this.azureEnvironment = azureEnvironment;
     }
@@ -36,7 +51,9 @@ public String getTenantId() {
         return tenantId;
     }
 
-    public CosmosAzureEnvironment getAzureEnvironment() {
-        return azureEnvironment;
+    public String getAuthEndpoint() {
+        String defaultAuthEndpoint = ACTIVE_DIRECTORY_ENDPOINT_MAP.get(azureEnvironment);
+        String authEndpoint = StringUtils.isNotEmpty(authEndpointOverride) ? authEndpointOverride : defaultAuthEndpoint;
+        return authEndpoint.replaceAll("/$", "") + "/";
     }
 }

sdk/cosmos/azure-cosmos-kafka-connect/src/main/java/com/azure/cosmos/kafka/connect/implementation/CosmosClientStore.java

@@ -14,20 +14,8 @@
 import com.azure.identity.ClientSecretCredentialBuilder;
 
 import java.time.Duration;
-import java.util.HashMap;
-import java.util.Map;
 
 public class CosmosClientStore {
-    private static final Map<CosmosAzureEnvironment, String> ACTIVE_DIRECTORY_ENDPOINT_MAP;
-    static {
-        // for now we maintain a static list within the SDK these values do not change very frequently
-        ACTIVE_DIRECTORY_ENDPOINT_MAP = new HashMap<>();
-        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE, "https://login.microsoftonline.com/");
-        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_CHINA, "https://login.chinacloudapi.cn/");
-        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_US_GOVERNMENT, "https://login.microsoftonline.us/");
-        ACTIVE_DIRECTORY_ENDPOINT_MAP.put(CosmosAzureEnvironment.AZURE_GERMANY, "https://login.microsoftonline.de/");
-    }
-
     public static CosmosAsyncClient getCosmosClient(
         CosmosAccountConfig accountConfig,
         String sourceName) {
@@ -59,10 +47,9 @@ public static CosmosAsyncClient getCosmosClient(
         if (accountConfig.getCosmosAuthConfig() instanceof CosmosMasterKeyAuthConfig) {
             cosmosClientBuilder.key(((CosmosMasterKeyAuthConfig) accountConfig.getCosmosAuthConfig()).getMasterKey());
         } else if (accountConfig.getCosmosAuthConfig() instanceof CosmosAadAuthConfig) {
-
             CosmosAadAuthConfig aadAuthConfig = (CosmosAadAuthConfig) accountConfig.getCosmosAuthConfig();
             ClientSecretCredential tokenCredential = new ClientSecretCredentialBuilder()
-                .authorityHost(ACTIVE_DIRECTORY_ENDPOINT_MAP.get(aadAuthConfig.getAzureEnvironment()).replaceAll("/$", "") + "/")
+                .authorityHost(aadAuthConfig.getAuthEndpoint())
                 .tenantId(aadAuthConfig.getTenantId())
                 .clientId(aadAuthConfig.getClientId())
                 .clientSecret(aadAuthConfig.getClientSecret())

sdk/cosmos/azure-cosmos-kafka-connect/src/main/java/com/azure/cosmos/kafka/connect/implementation/KafkaCosmosConfig.java

@@ -68,6 +68,14 @@ public class KafkaCosmosConfig extends AbstractConfig {
     private static final String AAD_CLIENT_SECRET_DISPLAY = "The client secret/password of the service principal.";
     private static final String DEFAULT_AAD_CLIENT_SECRET = Strings.Emtpy;
 
+    private static final String AAD_AUTH_ENDPOINT_OVERRIDE = "azure.cosmos.auth.aad.authEndpointOverride";
+    private static final String AAD_AUTH_ENDPOINT_OVERRIDE_DOC = "Overrides the Azure Active Directory (AAD) authentication endpoint. "
+        + "This is useful when the Cosmos DB account resides in a non-public Azure cloud environment such as Azure Air-Gapped Environment. "
+        + "By default, the SDK uses the standard AAD endpoint for the public Azure cloud. Set this value if your deployment requires a custom authority URI.";
+    private static final String AAD_AUTH_ENDPOINT_OVERRIDE_DISPLAY =
+        "The Azure Active Directory (AAD) authentication endpoint override. Set this if you are using a cloud environment other than public Azure.";
+    private static final String DEFAULT_AAD_AUTH_ENDPOINT_OVERRIDE = Strings.Emtpy;
+
     private static final String USE_GATEWAY_MODE = "azure.cosmos.mode.gateway";
     private static final String USE_GATEWAY_MODE_DOC = "Flag to indicate whether to use gateway mode. By default it is false, means SDK uses direct mode. https://learn.microsoft.com/azure/cosmos-db/nosql/sdk-connection-modes";
     private static final String USE_GATEWAY_MODE_DISPLAY = "Use gateway mode.";
@@ -206,6 +214,7 @@ private CosmosAccountConfig parseAccountConfig() {
             ACCOUNT_KEY,
             AAD_CLIENT_ID,
             AAD_CLIENT_SECRET,
+            AAD_AUTH_ENDPOINT_OVERRIDE,
             APPLICATION_NAME,
             USE_GATEWAY_MODE,
             PREFERRED_REGIONS_LIST);
@@ -219,6 +228,7 @@ private CosmosAccountConfig parseAccountConfigCore(
         String accountKeyConfig,
         String clientIdConfig,
         String clientSecretConfig,
+        String authEndpointOverrideConfig,
         String applicationNameConfig,
         String useGatewayModeConfig,
         String preferredRegionListConfig) {
@@ -230,7 +240,8 @@ private CosmosAccountConfig parseAccountConfigCore(
         String masterKey = this.getPassword(accountKeyConfig).value();
         String clientId = this.getString(clientIdConfig);
         String clientSecret = this.getPassword(clientSecretConfig).value();
-        CosmosAuthConfig authConfig = getAuthConfig(azureEnvironment, tenantId, authType, masterKey, clientId, clientSecret);
+        String authEndpointOverride = this.getString(authEndpointOverrideConfig);
+        CosmosAuthConfig authConfig = getAuthConfig(azureEnvironment, tenantId, authType, masterKey, clientId, clientSecret, authEndpointOverride);
 
         String applicationName = this.getString(applicationNameConfig);
         boolean useGatewayMode = this.getBoolean(useGatewayModeConfig);
@@ -250,13 +261,13 @@ private CosmosAuthConfig getAuthConfig(
         CosmosAuthType authType,
         String masterKey,
         String clientId,
-        String clientSecret) {
-
+        String clientSecret,
+        String authEndpointOverride) {
         switch (authType) {
             case MASTER_KEY:
                 return new CosmosMasterKeyAuthConfig(masterKey);
             case SERVICE_PRINCIPAL:
-                return new CosmosAadAuthConfig(clientId, clientSecret, tenantId, azureEnvironment);
+                return new CosmosAadAuthConfig(clientId, clientSecret, authEndpointOverride, tenantId, azureEnvironment);
             default:
                 throw new IllegalArgumentException("AuthType " + authType + " is not supported");
         }
@@ -286,6 +297,7 @@ private CosmosThroughputControlConfig parseThroughputControlConfig() {
                 THROUGHPUT_CONTROL_ACCOUNT_KEY,
                 THROUGHPUT_CONTROL_AAD_CLIENT_ID,
                 THROUGHPUT_CONTROL_AAD_CLIENT_SECRET,
+                AAD_AUTH_ENDPOINT_OVERRIDE,
                 APPLICATION_NAME,
                 THROUGHPUT_CONTROL_USE_GATEWAY_MODE,
                 THROUGHPUT_CONTROL_PREFERRED_REGIONS_LIST);
@@ -417,6 +429,17 @@ private static void defineAccountConfig(ConfigDef result) {
                 ConfigDef.Width.MEDIUM,
                 AAD_CLIENT_SECRET_DISPLAY
             )
+            .define(
+                AAD_AUTH_ENDPOINT_OVERRIDE,
+                ConfigDef.Type.STRING,
+                DEFAULT_AAD_AUTH_ENDPOINT_OVERRIDE,
+                ConfigDef.Importance.MEDIUM,
+                AAD_AUTH_ENDPOINT_OVERRIDE_DOC,
+                accountGroupName,
+                accountGroupOrder++,
+                ConfigDef.Width.LONG,
+                AAD_AUTH_ENDPOINT_OVERRIDE_DISPLAY
+            )
             .define(
                 APPLICATION_NAME,
                 ConfigDef.Type.STRING,
@@ -738,7 +761,8 @@ public static void validateThroughputControlConfig(Map<String, ConfigValue> conf
                 THROUGHPUT_CONTROL_AUTH_TYPE,
                 THROUGHPUT_CONTROL_ACCOUNT_KEY,
                 THROUGHPUT_CONTROL_AAD_CLIENT_ID,
-                THROUGHPUT_CONTROL_AAD_CLIENT_SECRET);
+                THROUGHPUT_CONTROL_AAD_CLIENT_SECRET,
+                AAD_AUTH_ENDPOINT_OVERRIDE);
         }
 
         // if throughput control is using aad auth, then only targetThroughput is supported
@@ -763,7 +787,8 @@ public static void validateCosmosAccountAuthConfig(Map<String, ConfigValue> conf
             AUTH_TYPE,
             ACCOUNT_KEY,
             AAD_CLIENT_ID,
-            AAD_CLIENT_SECRET);
+            AAD_CLIENT_SECRET,
+            AAD_AUTH_ENDPOINT_OVERRIDE);
     }
 
     public static void validateAccountAuthConfigCore(
@@ -772,7 +797,8 @@ public static void validateAccountAuthConfigCore(
         String authTypeConfig,
         String accountKeyConfig,
         String clientIdConfig,
-        String clientSecretConfig) {
+        String clientSecretConfig,
+        String authEndpointOverrideConfig) {
 
         CosmosAuthType authType = CosmosAuthType.fromName(configValueMap.get(authTypeConfig).value().toString());
         switch (authType) {
@@ -805,6 +831,18 @@ public static void validateAccountAuthConfigCore(
                         .get(clientSecretConfig)
                         .addErrorMessage("ClientSecret is required for Service Principal auth type");
                 }
+
+                String authEndpointOverride = configValueMap.get(authEndpointOverrideConfig).value().toString();
+                if (StringUtils.isNotEmpty(authEndpointOverride)) {
+                    try {
+                        new URL(authEndpointOverride);
+                    } catch (MalformedURLException e) {
+                        configValueMap
+                            .get(authEndpointOverrideConfig)
+                            .addErrorMessage("AuthEndpointOverride need to be valid URI format for Service Principal auth type");
+                    }
+                }
+
                 break;
             default:
                 throw new IllegalArgumentException("AuthType " + authType + " is not supported");