Skip to content

Commit e758006

Browse files
ibrandesCopilot
ibrandes
and
Copilot
authored
Storage - Improving Perf of StorageBearerTokenChallengeAuthorizationPolicy (Azure#46685)
* removing override that forces token fetch on initial request * scope refactoring * removing unecessary elements in signatures * refactoring authorizeRequestOnChallenge overrides * refactoring and strengthening parsing logic * some tests * fixing failing tests * adding overrides back to prevent breaking change * Update sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageBearerTokenChallengeAuthorizationPolicy.java Co-authored-by: Copilot <[email protected]> * Update sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageBearerTokenChallengeAuthorizationPolicyTests.java Co-authored-by: Copilot <[email protected]> * using CoreUtils.parseAuthenticateHeader instead of custom util method * adding changelog entry * adding scope and tenant verification * wip * added comment about failure * removing warning * adding custom parsing back and adjusting tests * many many tests for the parsing logic * undoing new parsing changes * moving stuff around * adjusting tests * more test adjustment * formatting --------- Co-authored-by: Copilot <[email protected]>
1 parent 54d16e6 commit e758006

File tree

3 files changed

+261
-201
lines changed

3 files changed

+261
-201
lines changed

‎sdk/storage/azure-storage-common/CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
### Breaking Changes
88

99
### Bugs Fixed
10+
- Addressed a performance regression introduced in version 12.27.0+ of the Azure Storage SDK for Java, where token
11+
retrieval became thread-local, causing a consistent 4 to 5-second delay across threads during initial authorization.
1012

1113
### Other Changes
1214

‎sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageBearerTokenChallengeAuthorizationPolicy.java‎

Lines changed: 57 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,20 @@
1818
import java.util.Collections;
1919
import java.util.HashMap;
2020
import java.util.Map;
21-
import java.util.Locale;
2221

2322
/**
24-
* The storage authorization policy which supports challenge.
23+
* The storage authorization policy which supports challenges.
2524
*/
2625
public class StorageBearerTokenChallengeAuthorizationPolicy extends BearerTokenAuthenticationPolicy {
27-
2826
private static final ClientLogger LOGGER = new ClientLogger(StorageBearerTokenChallengeAuthorizationPolicy.class);
2927

3028
private static final String DEFAULT_SCOPE = "/.default";
31-
static final String BEARER_TOKEN_PREFIX = "Bearer ";
29+
private static final String BEARER_TOKEN_PREFIX = "Bearer";
30+
private static final String RESOURCE_ID = "resource_id";
31+
private static final String AUTHORIZATION_URI = "authorization_uri";
3232

33-
private String[] scopes;
33+
// Immutable constructor scopes (base class retains them); challenge may supply new scopes dynamically.
34+
private final String[] initialScopes;
3435

3536
/**
3637
* Creates StorageBearerTokenChallengeAuthorizationPolicy.
@@ -40,56 +41,45 @@ public class StorageBearerTokenChallengeAuthorizationPolicy extends BearerTokenA
4041
*/
4142
public StorageBearerTokenChallengeAuthorizationPolicy(TokenCredential credential, String... scopes) {
4243
super(credential, scopes);
43-
this.scopes = scopes;
44+
this.initialScopes = CoreUtils.clone(scopes);
4445
}
4546

4647
@Override
4748
public Mono<Void> authorizeRequest(HttpPipelineCallContext context) {
48-
String[] scopes = this.scopes;
49-
scopes = getScopes(context, scopes);
50-
if (scopes == null) {
51-
return Mono.empty();
52-
} else {
53-
return setAuthorizationHeader(context, new TokenRequestContext().addScopes(scopes));
54-
}
49+
// Delegate to superclass to maintain previous behavior
50+
return super.authorizeRequest(context);
5551
}
5652

5753
@Override
5854
public void authorizeRequestSync(HttpPipelineCallContext context) {
59-
String[] scopes = this.scopes;
60-
scopes = getScopes(context, scopes);
61-
62-
if (scopes != null) {
63-
setAuthorizationHeaderSync(context, new TokenRequestContext().addScopes(scopes));
64-
}
55+
// Delegate to superclass to maintain previous behavior
56+
super.authorizeRequestSync(context);
6557
}
6658

6759
@Override
6860
public Mono<Boolean> authorizeRequestOnChallenge(HttpPipelineCallContext context, HttpResponse response) {
6961
String authHeader = response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE);
70-
Map<String, String> challenges = extractChallengeAttributes(authHeader, BEARER_TOKEN_PREFIX);
71-
72-
String scope = getScopeFromChallenges(challenges);
73-
String authorization = getAuthorizationFromChallenges(challenges);
62+
Map<String, String> attributes = extractChallengeAttributes(authHeader);
7463

75-
if (scope != null) {
76-
scope += DEFAULT_SCOPE;
77-
scopes = new String[] { scope };
78-
scopes = getScopes(context, scopes);
64+
if (attributes.isEmpty()) {
65+
return Mono.just(false);
7966
}
8067

81-
if (authorization != null) {
82-
String tenantId = extractTenantIdFromUri(authorization);
83-
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopes).setTenantId(tenantId);
84-
return setAuthorizationHeader(context, tokenRequestContext).thenReturn(true);
85-
}
68+
String resource = attributes.get(RESOURCE_ID);
69+
String authUrl = attributes.get(AUTHORIZATION_URI);
8670

87-
if (scope != null) {
88-
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopes);
89-
return setAuthorizationHeader(context, tokenRequestContext).thenReturn(true);
71+
// Create new scopes array to avoid mutating the original
72+
String[] scopesToUse = initialScopes;
73+
if (!CoreUtils.isNullOrEmpty(resource)) {
74+
scopesToUse = new String[] { resource.endsWith(DEFAULT_SCOPE) ? resource : resource + DEFAULT_SCOPE };
9075
}
9176

92-
return Mono.just(false);
77+
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopesToUse).setCaeEnabled(true);
78+
79+
if (!CoreUtils.isNullOrEmpty(authUrl)) {
80+
tokenRequestContext.setTenantId(extractTenantIdFromUri(authUrl));
81+
}
82+
return setAuthorizationHeader(context, tokenRequestContext).thenReturn(true);
9383
}
9484

9585
String extractTenantIdFromUri(String uri) {
@@ -108,66 +98,60 @@ String extractTenantIdFromUri(String uri) {
10898
@Override
10999
public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) {
110100
String authHeader = response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE);
111-
Map<String, String> challenges = extractChallengeAttributes(authHeader, BEARER_TOKEN_PREFIX);
112-
113-
String scope = getScopeFromChallenges(challenges);
114-
String authorization = getAuthorizationFromChallenges(challenges);
101+
Map<String, String> attributes = extractChallengeAttributes(authHeader);
115102

116-
if (scope != null) {
117-
scope += DEFAULT_SCOPE;
118-
scopes = new String[] { scope };
119-
scopes = getScopes(context, scopes);
103+
if (attributes.isEmpty()) {
104+
return false;
120105
}
121106

122-
if (authorization != null) {
123-
String tenantId = extractTenantIdFromUri(authorization);
124-
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopes).setTenantId(tenantId);
125-
setAuthorizationHeaderSync(context, tokenRequestContext);
126-
return true;
127-
}
107+
String resource = attributes.get(RESOURCE_ID);
108+
String authUrl = attributes.get(AUTHORIZATION_URI);
128109

129-
if (scope != null) {
130-
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopes);
131-
setAuthorizationHeaderSync(context, tokenRequestContext);
132-
return true;
110+
// Create new scopes array to avoid mutating the original
111+
String[] scopesToUse = initialScopes;
112+
if (!CoreUtils.isNullOrEmpty(resource)) {
113+
scopesToUse = new String[] { resource.endsWith(DEFAULT_SCOPE) ? resource : resource + DEFAULT_SCOPE };
133114
}
134115

135-
return false;
136-
}
116+
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(scopesToUse).setCaeEnabled(true);
137117

138-
String[] getScopes(HttpPipelineCallContext context, String[] scopes) {
139-
return CoreUtils.clone(scopes);
118+
if (!CoreUtils.isNullOrEmpty(authUrl)) {
119+
tokenRequestContext.setTenantId(extractTenantIdFromUri(authUrl));
120+
}
121+
122+
setAuthorizationHeaderSync(context, tokenRequestContext);
123+
return true;
140124
}
141125

142-
Map<String, String> extractChallengeAttributes(String header, String authChallengePrefix) {
143-
if (!isBearerChallenge(header, authChallengePrefix)) {
126+
static Map<String, String> extractChallengeAttributes(String header) {
127+
if (!isBearerChallenge(header)) {
144128
return Collections.emptyMap();
145129
}
146130

147-
header = header.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), "");
131+
// Don't lowercase the entire header as values can be corrupted. Also don't mutate the original header.
132+
String remainingHeader = header.substring(BEARER_TOKEN_PREFIX.length());
148133

149-
String[] attributes = header.split(" ");
134+
// Split on commas first; if no commas present fall back to spaces.
135+
String[] attributes = remainingHeader.contains(",") ? remainingHeader.split(",") : remainingHeader.split(" ");
150136
Map<String, String> attributeMap = new HashMap<>();
151137

152138
for (String pair : attributes) {
139+
// Skip empty strings and any strings that don't contain '='.
140+
// Allows scenarios where there is only an auth URI and no resource ID.
141+
if (CoreUtils.isNullOrEmpty(pair) || !pair.contains("=")) {
142+
continue;
143+
}
153144
String[] keyValue = pair.split("=");
154145

155-
attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", ""));
146+
// Support spaces after commas by trimming.
147+
attributeMap.put(keyValue[0].replaceAll("\"", "").trim(), keyValue[1].replaceAll("\"", "").trim());
156148
}
157149

158150
return attributeMap;
159151
}
160152

161-
static boolean isBearerChallenge(String authenticateHeader, String authChallengePrefix) {
153+
static boolean isBearerChallenge(String authenticateHeader) {
162154
return (!CoreUtils.isNullOrEmpty(authenticateHeader)
163-
&& authenticateHeader.toLowerCase(Locale.ROOT).startsWith(authChallengePrefix.toLowerCase(Locale.ROOT)));
164-
}
165-
166-
String getScopeFromChallenges(Map<String, String> challenges) {
167-
return challenges.get("resource_id");
168-
}
169-
170-
String getAuthorizationFromChallenges(Map<String, String> challenges) {
171-
return challenges.get("authorization_uri");
155+
&& authenticateHeader.regionMatches(true, 0, BEARER_TOKEN_PREFIX, 0, BEARER_TOKEN_PREFIX.length()));
172156
}
173157
}

0 commit comments

Comments
 (0)