diff --git a/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml b/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml
index ecf8041030041..c1225d64e03a3 100755
--- a/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml
+++ b/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml
@@ -298,10 +298,6 @@ the main ServiceBusClientBuilder. -->
-
-
-
-
diff --git a/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/KeyVaultAccessControlAsyncClient.java b/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/KeyVaultAccessControlAsyncClient.java
index 21d77ca75b7a9..d8f5e8eb2d1f4 100644
--- a/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/KeyVaultAccessControlAsyncClient.java
+++ b/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/KeyVaultAccessControlAsyncClient.java
@@ -384,10 +384,10 @@ Mono> setRoleDefinitionWithResponse(SetRoleDefi
options.getRoleDefinitionName(), parameters,
context.addData(AZ_TRACING_NAMESPACE_KEY, KEYVAULT_TRACING_NAMESPACE_VALUE))
.doOnRequest(ignored ->
- logger.verbose("Creating role assignment - {}", options.getRoleDefinitionName()))
- .doOnSuccess(response -> logger.verbose("Created role assignment - {}", response.getValue().getName()))
+ logger.verbose("Creating role definition - {}", options.getRoleDefinitionName()))
+ .doOnSuccess(response -> logger.verbose("Created role definition - {}", response.getValue().getName()))
.doOnError(error ->
- logger.warning("Failed to create role assignment - {}", options.getRoleDefinitionName(), error))
+ logger.warning("Failed to create role definition - {}", options.getRoleDefinitionName(), error))
.onErrorMap(KeyVaultAdministrationUtils::mapThrowableToKeyVaultAdministrationException)
.map(KeyVaultAccessControlAsyncClient::transformRoleDefinitionResponse);
} catch (RuntimeException e) {
@@ -471,11 +471,11 @@ Mono> getRoleDefinitionWithResponse(KeyVaultRol
return clientImpl.getRoleDefinitions()
.getWithResponseAsync(vaultUrl, roleScope.toString(), roleDefinitionName,
context.addData(AZ_TRACING_NAMESPACE_KEY, KEYVAULT_TRACING_NAMESPACE_VALUE))
- .doOnRequest(ignored -> logger.verbose("Retrieving role assignment - {}", roleDefinitionName))
+ .doOnRequest(ignored -> logger.verbose("Retrieving role definition - {}", roleDefinitionName))
.doOnSuccess(response ->
- logger.verbose("Retrieved role assignment - {}", response.getValue().getName()))
+ logger.verbose("Retrieved role definition - {}", response.getValue().getName()))
.doOnError(error ->
- logger.warning("Failed to retrieved role assignment - {}", roleDefinitionName, error))
+ logger.warning("Failed to retrieved role definition - {}", roleDefinitionName, error))
.onErrorMap(KeyVaultAdministrationUtils::mapThrowableToKeyVaultAdministrationException)
.map(KeyVaultAccessControlAsyncClient::transformRoleDefinitionResponse);
} catch (RuntimeException e) {
@@ -555,9 +555,9 @@ Mono> deleteRoleDefinitionWithResponse(KeyVaultRoleScope roleScop
return clientImpl.getRoleDefinitions()
.deleteWithResponseAsync(vaultUrl, roleScope.toString(), roleDefinitionName,
context.addData(AZ_TRACING_NAMESPACE_KEY, KEYVAULT_TRACING_NAMESPACE_VALUE))
- .doOnRequest(ignored -> logger.verbose("Deleting role assignment - {}", roleDefinitionName))
- .doOnSuccess(response -> logger.verbose("Deleted role assignment - {}", response.getValue().getName()))
- .doOnError(error -> logger.warning("Failed to delete role assignment - {}", roleDefinitionName, error))
+ .doOnRequest(ignored -> logger.verbose("Deleting role definition - {}", roleDefinitionName))
+ .doOnSuccess(response -> logger.verbose("Deleted role definition - {}", response.getValue().getName()))
+ .doOnError(error -> logger.warning("Failed to delete role definition - {}", roleDefinitionName, error))
.onErrorMap(KeyVaultAdministrationUtils::mapThrowableToKeyVaultAdministrationException)
.map(response -> (Response) new SimpleResponse(response, null))
.onErrorResume(KeyVaultAdministrationException.class, e ->
@@ -897,7 +897,7 @@ Mono> getRoleAssignmentWithResponse(KeyVaultRol
.doOnSuccess(response ->
logger.verbose("Retrieved role assignment - {}", response.getValue().getName()))
.doOnError(error ->
- logger.warning("Failed to retrieved role assignment - {}", roleAssignmentName, error))
+ logger.warning("Failed to retrieve role assignment - {}", roleAssignmentName, error))
.onErrorMap(KeyVaultAdministrationUtils::mapThrowableToKeyVaultAdministrationException)
.map(KeyVaultAccessControlAsyncClient::transformRoleAssignmentResponse);
} catch (RuntimeException e) {
diff --git a/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/implementation/KeyVaultCredentialPolicy.java b/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/implementation/KeyVaultCredentialPolicy.java
index 947721a0cdc4a..25140f45b35d0 100644
--- a/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/implementation/KeyVaultCredentialPolicy.java
+++ b/sdk/keyvault/azure-security-keyvault-administration/src/main/java/com/azure/security/keyvault/administration/implementation/KeyVaultCredentialPolicy.java
@@ -1,125 +1,186 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
-
package com.azure.security.keyvault.administration.implementation;
import com.azure.core.credential.TokenCredential;
import com.azure.core.credential.TokenRequestContext;
import com.azure.core.http.HttpPipelineCallContext;
-import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpRequest;
import com.azure.core.http.HttpResponse;
-import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.http.policy.BearerTokenAuthenticationPolicy;
import com.azure.core.util.CoreUtils;
-import com.azure.core.util.logging.ClientLogger;
+import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.util.Collections;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
-import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
/**
- * A policy that authenticates requests with Azure Key Vault service. The content added by this policy is leveraged
- * in {@link TokenCredential} to get and set the correct "Authorization" header value.
+ * A policy that authenticates requests with the Azure Key Vault service. The content added by this policy is
+ * leveraged in {@link TokenCredential} to get and set the correct "Authorization" header value.
*
* @see TokenCredential
*/
-public final class KeyVaultCredentialPolicy implements HttpPipelinePolicy {
- private final ClientLogger logger = new ClientLogger(KeyVaultCredentialPolicy.class);
- private static final String WWW_AUTHENTICATE = "WWW-Authenticate";
+public class KeyVaultCredentialPolicy extends BearerTokenAuthenticationPolicy {
private static final String BEARER_TOKEN_PREFIX = "Bearer ";
- private static final String AUTHORIZATION = "Authorization";
- private final ScopeTokenCache cache;
+ private static final String CONTENT_LENGTH_HEADER = "Content-Length";
+ private static final String KEY_VAULT_STASHED_CONTENT_KEY = "KeyVaultCredentialPolicyStashedBody";
+ private static final String KEY_VAULT_STASHED_CONTENT_LENGTH_KEY = "KeyVaultCredentialPolicyStashedContentLength";
+ private static final String WWW_AUTHENTICATE = "WWW-Authenticate";
+ private static final ConcurrentMap SCOPE_CACHE = new ConcurrentHashMap<>();
+ private String scope;
/**
- * Creates KeyVaultCredentialPolicy.
+ * Creates a {@link KeyVaultCredentialPolicy}.
*
- * @param credential the token credential to authenticate the request
+ * @param credential The token credential to authenticate the request.
*/
public KeyVaultCredentialPolicy(TokenCredential credential) {
- Objects.requireNonNull(credential);
-
- this.cache = new ScopeTokenCache(credential::getToken);
- }
-
- /**
- * Adds the required header to authenticate a request to Azure Key Vault service.
- *
- * @param context The request {@link HttpPipelineCallContext context}.
- * @param next The next HTTP pipeline policy to process the {@link HttpPipelineCallContext context's} request
- * after this policy completes.
- * @return A {@link Mono} representing the {@link HttpResponse HTTP response} that will arrive asynchronously.
- */
- @Override
- public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
- if (!context.getHttpRequest().getUrl().getProtocol().startsWith("https")) {
- return Mono.error(new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme"));
- }
-
- return next.clone().process()
- .doOnNext(httpResponse -> {
- // KV follows challenge based auth. Currently every service
- // call hit the endpoint for challenge and then resend the
- // request with token. The challenge response body is not
- // consumed, not draining/closing the body will result in leak.
- // Ref: https://github.com/Azure/azure-sdk-for-java/issues/7934
- // https://github.com/Azure/azure-sdk-for-java/issues/10467
- try {
- httpResponse.getBody().subscribe().dispose();
- } catch (RuntimeException ignored) {
- logger.logExceptionAsWarning(ignored);
- }
- // The ReactorNettyHttpResponse::close() should be sufficient
- // and should take care similar body disposal but looks like that
- // is not happening, need to re-visit the close() method.
- })
- .map(res -> res.getHeaderValue(WWW_AUTHENTICATE))
- .map(header -> extractChallenge(header, BEARER_TOKEN_PREFIX))
- .flatMap(map -> {
- cache.setTokenRequest(new TokenRequestContext().addScopes(map.get("resource") + "/.default"));
- return cache.getToken();
- })
- .flatMap(token -> {
- context.getHttpRequest().setHeader(AUTHORIZATION, BEARER_TOKEN_PREFIX + token.getToken());
- return next.process();
- });
+ super(credential);
}
/**
- * Extracts the challenge off the authentication header.
+ * Extracts attributes off the bearer challenge in the authentication header.
*
- * @param authenticateHeader The authentication header containing all the challenges.
+ * @param authenticateHeader The authentication header containing the challenge.
* @param authChallengePrefix The authentication challenge name.
- * @return A challenge map.
+ *
+ * @return A challenge attributes map.
*/
- private static Map extractChallenge(String authenticateHeader, String authChallengePrefix) {
- if (!isValidChallenge(authenticateHeader, authChallengePrefix)) {
- return null;
+ private static Map extractChallengeAttributes(String authenticateHeader,
+ String authChallengePrefix) {
+ if (!isBearerChallenge(authenticateHeader, authChallengePrefix)) {
+ return Collections.emptyMap();
}
authenticateHeader =
authenticateHeader.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), "");
- String[] challenges = authenticateHeader.split(", ");
- Map challengeMap = new HashMap<>();
+ String[] attributes = authenticateHeader.split(", ");
+ Map attributeMap = new HashMap<>();
- for (String pair : challenges) {
+ for (String pair : attributes) {
String[] keyValue = pair.split("=");
- challengeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", ""));
+
+ attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", ""));
}
- return challengeMap;
+ return attributeMap;
}
/**
* Verifies whether a challenge is bearer or not.
*
- * @param authenticateHeader The authentication header containing all the challenges.
+ * @param authenticateHeader The authentication header containing all the challenges.
* @param authChallengePrefix The authentication challenge name.
- * @return A boolean indicating tha challenge is valid or not.
+ * @return A boolean indicating if the challenge is a bearer challenge or not.
*/
- private static boolean isValidChallenge(String authenticateHeader, String authChallengePrefix) {
+ private static boolean isBearerChallenge(String authenticateHeader, String authChallengePrefix) {
return (!CoreUtils.isNullOrEmpty(authenticateHeader)
&& authenticateHeader.toLowerCase(Locale.ROOT).startsWith(authChallengePrefix.toLowerCase(Locale.ROOT)));
}
+
+ @Override
+ public Mono authorizeRequest(HttpPipelineCallContext context) {
+ return Mono.defer(() -> {
+ HttpRequest request = context.getHttpRequest();
+
+ // If this policy doesn't have an authorityScope cached try to get it from the static challenge cache.
+ if (this.scope == null) {
+ String authority = getRequestAuthority(request);
+ this.scope = SCOPE_CACHE.get(authority);
+ }
+
+ if (this.scope != null) {
+ // We fetched the scope from the cache, but we have not initialized the scopes in the base yet.
+ TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes(this.scope);
+
+ return setAuthorizationHeader(context, tokenRequestContext);
+ }
+
+ // The body is removed from the initial request because Key Vault supports other authentication schemes which
+ // also protect the body of the request. As a result, before we know the auth scheme we need to avoid sending
+ // an unprotected body to Key Vault. We don't currently support this enhanced auth scheme in the SDK but we
+ // still don't want to send any unprotected data to vaults which require it.
+
+ // Do not overwrite previous contents if retrying after initial request failed (e.g. timeout).
+ if (!context.getData(KEY_VAULT_STASHED_CONTENT_KEY).isPresent()) {
+ if (request.getBody() != null) {
+ context.setData(KEY_VAULT_STASHED_CONTENT_KEY, request.getBody());
+ context.setData(KEY_VAULT_STASHED_CONTENT_LENGTH_KEY,
+ request.getHeaders().getValue(CONTENT_LENGTH_HEADER));
+ request.setHeader(CONTENT_LENGTH_HEADER, "0");
+ request.setBody((Flux) null);
+ }
+ }
+
+ return Mono.empty();
+ });
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context, HttpResponse response) {
+ return Mono.defer(() -> {
+ HttpRequest request = context.getHttpRequest();
+ Optional