Skip to content

Commit 5bb484f

Browse files
milismsftmilismsft
authored
Cosmos: preview for AAD support (#12622)
* Add initial implementation to pass an AAD token to the backend. * Address PR comments * Add AAD authorization test against Cosmos public emulator. Add implementation for missed cases where authorization token migt be computed. * update pom related dependency * update pom dependency * test updates * address PR feedback * address PR feedback * Bug fixes. * enable AAD auth in the Cosmos public emulator * update Cosmos emulator startup switch * update test case to separate access via different clients * Address PR feedback. * Remove constructor which creates unused Cosmos resources. * use HOST and MASTER_KEY for Cosmos connections; these will default to Cosmos public emulator settings. * Update test case expectations. * update Sping related test expectations. * Update Spring tests expectations and fix couple error cases when passing empty strings for endpoints and master keys. * Fix for scope resolution * comment out the test until the CI only failure running public emulator is understood. * update POM dependencies. * Fix merge related issue. * various fixes related to copy/clone of an existing Cosmos client instance. * update test to account for null values such as key, endpoint or credential properties.
1 parent f5fa3a4 commit 5bb484f

File tree

21 files changed

+756
-158
lines changed

21 files changed

+756
-158
lines changed

‎eng/pipelines/templates/stages/cosmos-sdk-client.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ stages:
129129
PreRunSteps:
130130
- template: /eng/common/pipelines/templates/steps/cosmos-emulator.yml
131131
parameters:
132-
StartParameters: '-PartitionCount 50 -Consistency Strong -Timeout 600'
132+
StartParameters: '-EnableAadAuthentication -PartitionCount 50 -Consistency Strong -Timeout 600'
133133
- powershell: |
134134
$Key = 'C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=='
135135
$password = ConvertTo-SecureString -String $Key -Force -AsPlainText

‎sdk/cosmos/azure-cosmos/pom.xml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,26 @@ Licensed under the MIT License.
122122
<artifactId>azure-core</artifactId>
123123
<version>1.8.1</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
124124
</dependency>
125+
<dependency>
126+
<groupId>com.azure</groupId>
127+
<artifactId>azure-identity</artifactId>
128+
<version>1.1.2</version> <!-- {x-version-update;com.azure:azure-identity;dependency} -->
129+
<scope>test</scope>
130+
</dependency>
131+
132+
<dependency>
133+
<groupId>com.azure</groupId>
134+
<artifactId>azure-core-http-netty</artifactId>
135+
<version>1.6.1</version> <!-- {x-version-update;com.azure:azure-core-http-netty;dependency} -->
136+
<optional>true</optional>
137+
<exclusions>
138+
<exclusion>
139+
<groupId>com.azure</groupId>
140+
<artifactId>azure-core</artifactId>
141+
</exclusion>
142+
</exclusions>
143+
</dependency>
144+
125145
<dependency>
126146
<groupId>com.fasterxml.jackson.module</groupId>
127147
<artifactId>jackson-module-afterburner</artifactId>

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import com.azure.core.annotation.ServiceClient;
66
import com.azure.core.credential.AzureKeyCredential;
7+
import com.azure.core.credential.TokenCredential;
78
import com.azure.core.util.Context;
89
import com.azure.core.util.tracing.Tracer;
910
import com.azure.cosmos.implementation.AsyncDocumentClient;
@@ -55,6 +56,7 @@ public final class CosmosAsyncClient implements Closeable {
5556
private final List<CosmosPermissionProperties> permissions;
5657
private final CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
5758
private final AzureKeyCredential credential;
59+
private final TokenCredential tokenCredential;
5860
private final boolean sessionCapturingOverride;
5961
private final boolean enableTransportClientSharing;
6062
private final TracerProvider tracerProvider;
@@ -80,6 +82,7 @@ public final class CosmosAsyncClient implements Closeable {
8082
this.permissions = builder.getPermissions();
8183
this.cosmosAuthorizationTokenResolver = builder.getAuthorizationTokenResolver();
8284
this.credential = builder.getCredential();
85+
this.tokenCredential = builder.getTokenCredential();
8386
this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
8487
this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
8588
this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
@@ -95,6 +98,7 @@ public final class CosmosAsyncClient implements Closeable {
9598
.withCredential(this.credential)
9699
.withTransportClientSharing(this.enableTransportClientSharing)
97100
.withContentResponseOnWriteEnabled(this.contentResponseOnWriteEnabled)
101+
.withTokenCredential(this.tokenCredential)
98102
.build();
99103
}
100104

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosBridgeInternal.java‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import com.azure.cosmos.implementation.AsyncDocumentClient;
77
import com.azure.cosmos.implementation.ConnectionPolicy;
88
import com.azure.cosmos.implementation.Document;
9+
import com.azure.cosmos.implementation.Strings;
910
import com.azure.cosmos.implementation.Warning;
1011
import com.azure.cosmos.implementation.query.Transformer;
1112
import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -86,16 +87,38 @@ public static ConnectionPolicy getConnectionPolicy(CosmosClientBuilder cosmosCli
8687
@Warning(value = INTERNAL_USE_ONLY_WARNING)
8788
public static CosmosClientBuilder cloneCosmosClientBuilder(CosmosClientBuilder builder) {
8889
CosmosClientBuilder copy = new CosmosClientBuilder();
90+
if (!Strings.isNullOrEmpty(builder.getEndpoint())) {
91+
copy.endpoint(builder.getEndpoint());
92+
}
8993

90-
copy.endpoint(builder.getEndpoint())
91-
.key(builder.getKey())
94+
if (!Strings.isNullOrEmpty(builder.getKey())) {
95+
copy.key(builder.getKey());
96+
}
97+
98+
if (!Strings.isNullOrEmpty(builder.getResourceToken())) {
99+
copy.resourceToken(builder.getResourceToken());
100+
}
101+
102+
if (builder.getCredential() != null) {
103+
copy.credential(builder.getCredential());
104+
}
105+
106+
if (builder.getTokenCredential() != null) {
107+
copy.credential(builder.getTokenCredential());
108+
}
109+
110+
if (builder.getPermissions() != null) {
111+
copy.permissions(builder.getPermissions());
112+
}
113+
114+
if (builder.getAuthorizationTokenResolver() != null) {
115+
copy.authorizationTokenResolver(builder.getAuthorizationTokenResolver());
116+
}
117+
118+
copy
92119
.directMode(builder.getDirectConnectionConfig())
93120
.gatewayMode(builder.getGatewayConnectionConfig())
94121
.consistencyLevel(builder.getConsistencyLevel())
95-
.credential(builder.getCredential())
96-
.permissions(builder.getPermissions())
97-
.authorizationTokenResolver(builder.getAuthorizationTokenResolver())
98-
.resourceToken(builder.getResourceToken())
99122
.contentResponseOnWriteEnabled(builder.isContentResponseOnWriteEnabled())
100123
.userAgentSuffix(builder.getUserAgentSuffix())
101124
.throttlingRetryOptions(builder.getThrottlingRetryOptions())

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java‎

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import com.azure.core.annotation.ServiceClientBuilder;
66
import com.azure.core.credential.AzureKeyCredential;
7+
import com.azure.core.credential.TokenCredential;
78
import com.azure.cosmos.implementation.Configs;
89
import com.azure.cosmos.implementation.ConnectionPolicy;
910
import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
@@ -12,6 +13,7 @@
1213

1314
import java.util.Collections;
1415
import java.util.List;
16+
import java.util.Objects;
1517

1618
/**
1719
* Helper class to build CosmosAsyncClient {@link CosmosAsyncClient} and CosmosClient {@link CosmosClient}
@@ -80,6 +82,7 @@ public class CosmosClientBuilder {
8082
private Configs configs = new Configs();
8183
private String serviceEndpoint;
8284
private String keyOrResourceToken;
85+
private TokenCredential tokenCredential;
8386
private ConnectionPolicy connectionPolicy;
8487
private GatewayConnectionConfig gatewayConnectionConfig;
8588
private DirectConnectionConfig directConnectionConfig;
@@ -199,7 +202,12 @@ CosmosAuthorizationTokenResolver getAuthorizationTokenResolver() {
199202
*/
200203
CosmosClientBuilder authorizationTokenResolver(
201204
CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver) {
202-
this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
205+
this.cosmosAuthorizationTokenResolver = Objects.requireNonNull(cosmosAuthorizationTokenResolver,
206+
"'cosmosAuthorizationTokenResolver' cannot be null.");
207+
this.keyOrResourceToken = null;
208+
this.credential = null;
209+
this.permissions = null;
210+
this.tokenCredential = null;
203211
return this;
204212
}
205213

@@ -219,7 +227,7 @@ String getEndpoint() {
219227
* @return current Builder
220228
*/
221229
public CosmosClientBuilder endpoint(String endpoint) {
222-
this.serviceEndpoint = endpoint;
230+
this.serviceEndpoint = Objects.requireNonNull(endpoint, "'endpoint' cannot be null.");
223231
return this;
224232
}
225233

@@ -241,12 +249,16 @@ String getKey() {
241249
* @return current Builder.
242250
*/
243251
public CosmosClientBuilder key(String key) {
244-
this.keyOrResourceToken = key;
252+
this.keyOrResourceToken = Objects.requireNonNull(key, "'key' cannot be null.");
253+
this.cosmosAuthorizationTokenResolver = null;
254+
this.credential = null;
255+
this.permissions = null;
256+
this.tokenCredential = null;
245257
return this;
246258
}
247259

248260
/**
249-
* Sets a resource token used to perform authentication
261+
* Gets a resource token used to perform authentication
250262
* for accessing resource.
251263
*
252264
* @return the resourceToken
@@ -263,7 +275,37 @@ String getResourceToken() {
263275
* @return current Builder.
264276
*/
265277
public CosmosClientBuilder resourceToken(String resourceToken) {
266-
this.keyOrResourceToken = resourceToken;
278+
this.keyOrResourceToken = Objects.requireNonNull(resourceToken, "'resourceToken' cannot be null.");
279+
this.cosmosAuthorizationTokenResolver = null;
280+
this.credential = null;
281+
this.permissions = null;
282+
this.tokenCredential = null;
283+
return this;
284+
}
285+
286+
/**
287+
* Gets a token credential instance used to perform authentication
288+
* for accessing resource.
289+
*
290+
* @return the token credential.
291+
*/
292+
TokenCredential getTokenCredential() {
293+
return tokenCredential;
294+
}
295+
296+
/**
297+
* Sets the {@link TokenCredential} used to authorize requests sent to the service.
298+
*
299+
* @param credential {@link TokenCredential}.
300+
* @return the updated CosmosClientBuilder
301+
* @throws NullPointerException If {@code credential} is {@code null}.
302+
*/
303+
public CosmosClientBuilder credential(TokenCredential credential) {
304+
this.tokenCredential = Objects.requireNonNull(credential, "'credential' cannot be null.");
305+
this.keyOrResourceToken = null;
306+
this.cosmosAuthorizationTokenResolver = null;
307+
this.credential = null;
308+
this.permissions = null;
267309
return this;
268310
}
269311

@@ -285,7 +327,11 @@ List<CosmosPermissionProperties> getPermissions() {
285327
* @return current Builder.
286328
*/
287329
public CosmosClientBuilder permissions(List<CosmosPermissionProperties> permissions) {
288-
this.permissions = permissions;
330+
this.permissions = Objects.requireNonNull(permissions, "'permissions' cannot be null.");
331+
this.keyOrResourceToken = null;
332+
this.cosmosAuthorizationTokenResolver = null;
333+
this.credential = null;
334+
this.tokenCredential = null;
289335
return this;
290336
}
291337

@@ -338,7 +384,11 @@ AzureKeyCredential getCredential() {
338384
* @return current cosmosClientBuilder
339385
*/
340386
public CosmosClientBuilder credential(AzureKeyCredential credential) {
341-
this.credential = credential;
387+
this.credential = Objects.requireNonNull(credential, "'cosmosKeyCredential' cannot be null.");
388+
this.keyOrResourceToken = null;
389+
this.cosmosAuthorizationTokenResolver = null;
390+
this.permissions = null;
391+
this.tokenCredential = null;
342392
return this;
343393
}
344394

@@ -689,7 +739,8 @@ private void validateConfig() {
689739
ifThrowIllegalArgException(this.serviceEndpoint == null,
690740
"cannot buildAsyncClient client without service endpoint");
691741
ifThrowIllegalArgException(
692-
this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty()) && this.credential == null,
742+
this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty())
743+
&& this.credential == null && this.tokenCredential == null,
693744
"cannot buildAsyncClient client without any one of key, resource token, permissions, and "
694745
+ "azure key credential");
695746
ifThrowIllegalArgException(credential != null && StringUtils.isEmpty(credential.getKey()),

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AadTokenAuthorizationHelper.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
package com.azure.cosmos.implementation;
4+
5+
import com.azure.core.credential.SimpleTokenCache;
6+
import org.slf4j.Logger;
7+
import org.slf4j.LoggerFactory;
8+
import reactor.core.publisher.Mono;
9+
10+
import java.io.UnsupportedEncodingException;
11+
import java.net.URLEncoder;
12+
13+
/**
14+
* This class is used internally and act as a helper in authorization of
15+
* AAD tokens and its supporting method.
16+
*
17+
*/
18+
public class AadTokenAuthorizationHelper {
19+
public static final String AAD_AUTH_SCHEMA_TYPE_SEGMENT = "type";
20+
public static final String AAD_AUTH_VERSION_SEGMENT = "ver";
21+
public static final String AAD_AUTH_SIGNATURE_SEGMENT = "sig";
22+
public static final String AAD_AUTH_SCHEMA_TYPE_VALUE = "aad";
23+
public static final String AAD_AUTH_VERSION_VALUE = "1.0";
24+
public static final String AAD_AUTH_TOKEN_COSMOS_SCOPE = "https://cosmos.azure.com/.default";
25+
private static final String AUTH_PREFIX =
26+
AAD_AUTH_SCHEMA_TYPE_SEGMENT + "=" + AAD_AUTH_SCHEMA_TYPE_VALUE
27+
+ "&"
28+
+ AAD_AUTH_VERSION_SEGMENT + "=" + AAD_AUTH_VERSION_VALUE
29+
+ "&"
30+
+ AAD_AUTH_SIGNATURE_SEGMENT + "=";
31+
private static final Logger logger = LoggerFactory.getLogger(AadTokenAuthorizationHelper.class);
32+
33+
/**
34+
* This method will try to fetch the AAD token to access the resource and add it to the request headers.

35+
*
36+
* @param request the request headers.
37+
* @param simpleTokenCache token cache that supports caching a token and refreshing it.
38+
* @return the request headers with authorization header updated.
39+
*/
40+
public static Mono<RxDocumentServiceRequest> populateAuthorizationHeader(RxDocumentServiceRequest request, SimpleTokenCache simpleTokenCache) {
41+
if (request == null || request.getHeaders() == null) {
42+
throw new IllegalArgumentException("request");
43+
}
44+
if (simpleTokenCache == null) {
45+
throw new IllegalArgumentException("simpleTokenCache");
46+
}
47+
48+
return getAuthorizationToken(simpleTokenCache)
49+
.map(authorization -> {
50+
request.getHeaders().put(HttpConstants.HttpHeaders.AUTHORIZATION, authorization);
51+
return request;
52+
});
53+
}
54+
55+
public static Mono<String> getAuthorizationToken(SimpleTokenCache simpleTokenCache) {
56+
return simpleTokenCache.getToken()
57+
.map(accessToken -> {
58+
String authorization;
59+
String authorizationPayload = AUTH_PREFIX + accessToken.getToken();
60+
61+
try {
62+
authorization = URLEncoder.encode(authorizationPayload, "UTF-8");
63+
} catch (UnsupportedEncodingException e) {
64+
throw new IllegalStateException("Failed to encode authorization token.", e);
65+
}
66+
67+
return authorization;
68+
});
69+
}
70+
}

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
package com.azure.cosmos.implementation;
44

55
import com.azure.core.credential.AzureKeyCredential;
6+
import com.azure.core.credential.TokenCredential;
67
import com.azure.cosmos.ConsistencyLevel;
78
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
89
import com.azure.cosmos.models.CosmosItemIdentity;
@@ -72,6 +73,7 @@ class Builder {
7273
URI serviceEndpoint;
7374
CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
7475
AzureKeyCredential credential;
76+
TokenCredential tokenCredential;
7577
boolean sessionCapturingOverride;
7678
boolean transportClientSharing;
7779
boolean contentResponseOnWriteEnabled;
@@ -172,6 +174,17 @@ public Builder withTokenResolver(CosmosAuthorizationTokenResolver cosmosAuthoriz
172174
return this;
173175
}
174176

177+
/**
178+
* This method will accept functional interface TokenCredential which helps in generation authorization
179+
* token per request. AsyncDocumentClient can be successfully initialized with this API without passing any MasterKey, ResourceToken or PermissionFeed.
180+
* @param tokenCredential the token credential
181+
* @return current Builder.
182+
*/
183+
public Builder withTokenCredential(TokenCredential tokenCredential) {
184+
this.tokenCredential = tokenCredential;
185+
return this;
186+
}
187+
175188
private void ifThrowIllegalArgException(boolean value, String error) {
176189
if (value) {
177190
throw new IllegalArgumentException(error);
@@ -180,10 +193,10 @@ private void ifThrowIllegalArgException(boolean value, String error) {
180193

181194
public AsyncDocumentClient build() {
182195

183-
ifThrowIllegalArgException(this.serviceEndpoint == null, "cannot buildAsyncClient client without service endpoint");
196+
ifThrowIllegalArgException(this.serviceEndpoint == null || StringUtils.isEmpty(this.serviceEndpoint.toString()), "cannot buildAsyncClient client without service endpoint");
184197
ifThrowIllegalArgException(
185198
this.masterKeyOrResourceToken == null && (permissionFeed == null || permissionFeed.isEmpty())
186-
&& this.credential == null,
199+
&& this.credential == null && this.tokenCredential == null,
187200
"cannot buildAsyncClient client without any one of masterKey, " +
188201
"resource token, permissionFeed and azure key credential");
189202
ifThrowIllegalArgException(credential != null && StringUtils.isEmpty(credential.getKey()),
@@ -197,6 +210,7 @@ public AsyncDocumentClient build() {
197210
configs,
198211
cosmosAuthorizationTokenResolver,
199212
credential,
213+
tokenCredential,
200214
sessionCapturingOverride,
201215
transportClientSharing,
202216
contentResponseOnWriteEnabled);

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AuthorizationTokenType.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,5 +11,6 @@ public enum AuthorizationTokenType {
1111
SystemReadOnly,
1212
SystemReadWrite,
1313
SystemAll,
14-
ResourceToken
14+
ResourceToken,
15+
AadToken
1516
}

0 commit comments

Comments
 (0)