From 63a593c2e301a27824eb0879f42392d099f5a322 Mon Sep 17 00:00:00 2001 From: vcolin7 Date: Tue, 15 Oct 2024 15:02:57 -0700 Subject: [PATCH] Added code to handle claims in authentication challenges (#41814) --- .../CHANGELOG.md | 9 +- .../KeyVaultCredentialPolicy.java | 205 +++++- .../KeyVaultCredentialPolicyTest.java | 581 +++++++++++++++--- .../CHANGELOG.md | 3 + .../KeyVaultCredentialPolicy.java | 201 +++++- .../KeyVaultCredentialPolicyTest.java | 525 ++++++++++++++-- .../azure-security-keyvault-keys/CHANGELOG.md | 15 +- .../KeyVaultCredentialPolicy.java | 206 ++++++- .../keys/KeyVaultCredentialPolicyTest.java | 524 ++++++++++++++-- .../CHANGELOG.md | 8 - .../KeyVaultCredentialPolicy.java | 201 +++++- .../secrets/KeyVaultCredentialPolicyTest.java | 523 ++++++++++++++-- 12 files changed, 2668 insertions(+), 333 deletions(-) diff --git a/sdk/keyvault/azure-security-keyvault-administration/CHANGELOG.md b/sdk/keyvault/azure-security-keyvault-administration/CHANGELOG.md index a6579dcf930ad..c89440bc14849 100644 --- a/sdk/keyvault/azure-security-keyvault-administration/CHANGELOG.md +++ b/sdk/keyvault/azure-security-keyvault-administration/CHANGELOG.md @@ -2,6 +2,9 @@ ## 4.6.0 (2024-10-15) +## Features Added +- Added support for Continuous Access Evaluation (CAE). ([#41814](https://github.com/Azure/azure-sdk-for-java/pull/41814)) + ### Other Changes #### Dependency Updates @@ -19,7 +22,6 @@ - Upgraded `azure-core-http-netty` from `1.15.3` to version `1.15.4`. - Upgraded `azure-core` from `1.51.0` to version `1.52.0`. - ## 4.5.7 (2024-08-24) ### Other Changes @@ -29,7 +31,6 @@ - Upgraded `azure-core` from `1.50.0` to version `1.51.0`. - Upgraded `azure-core-http-netty` from `1.15.2` to version `1.15.3`. - ## 4.5.6 (2024-07-29) ### Other Changes @@ -49,7 +50,6 @@ - Upgraded `azure-core` from `1.49.0` to version `1.49.1`. - Upgraded `azure-core-http-netty` from `1.15.0` to version `1.15.1`. - ## 4.5.4 (2024-05-13) ### Other Changes @@ -74,7 +74,6 @@ - Upgraded `azure-core` from `1.47.0` to version `1.48.0`. - Upgraded `azure-core-http-netty` from `1.14.1` to version `1.14.2`. - ## 4.5.1 (2024-03-20) ### Other Changes @@ -84,7 +83,6 @@ - Upgraded `azure-core` from `1.46.0` to version `1.47.0`. - Upgraded `azure-core-http-netty` from `1.14.0` to version `1.14.1`. - ## 4.5.0 (2024-02-22) Changes when compared to the last stable release (`4.4.3`) include: @@ -172,7 +170,6 @@ Changes when compared to the last stable release (`4.4.3`) include: - Upgraded `azure-core` from `1.40.0` to version `1.41.0`. - Upgraded `azure-core-http-netty` from `1.13.4` to version `1.13.5`. - ## 4.3.3 (2023-06-20) ### Other Changes 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 79ca22b1d6913..8043cdc730710 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 @@ -4,11 +4,13 @@ import com.azure.core.credential.TokenCredential; import com.azure.core.credential.TokenRequestContext; -import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpPipelineCallContext; +import com.azure.core.http.HttpPipelineNextPolicy; +import com.azure.core.http.HttpPipelineNextSyncPolicy; import com.azure.core.http.HttpRequest; import com.azure.core.http.HttpResponse; import com.azure.core.http.policy.BearerTokenAuthenticationPolicy; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; @@ -19,6 +21,7 @@ import java.net.URISyntaxException; import java.net.URL; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashMap; import java.util.Locale; @@ -27,6 +30,9 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import static com.azure.core.http.HttpHeaderName.CONTENT_LENGTH; +import static com.azure.core.http.HttpHeaderName.WWW_AUTHENTICATE; + /** * 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. @@ -67,16 +73,17 @@ private static Map extractChallengeAttributes(String authenticat return Collections.emptyMap(); } - authenticateHeader = - authenticateHeader.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), ""); - - String[] attributes = authenticateHeader.split(", "); + String[] attributes = authenticateHeader + .replace("\"", "") + .substring(authChallengePrefix.length()) + .split(","); Map attributeMap = new HashMap<>(); for (String pair : attributes) { - String[] keyValue = pair.split("="); + // Using trim is ugly, but we need it here because currently the 'claims' attribute comes after two spaces. + String[] keyValue = pair.trim().split("=", 2); - attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", "")); + attributeMap.put(keyValue[0], keyValue[1]); } return attributeMap; @@ -102,31 +109,31 @@ public Mono authorizeRequest(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); 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. + // 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(HttpHeaderName.CONTENT_LENGTH)); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((Flux) null); } } @@ -145,13 +152,12 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((Flux) contentOptional.get()); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); Map challengeAttributes = - extractChallengeAttributes(response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE), - BEARER_TOKEN_PREFIX); + extractChallengeAttributes(response.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); String scope = challengeAttributes.get("resource"); if (scope != null) { @@ -201,7 +207,23 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims( + new String(Base64Util.decodeString(claims), StandardCharsets.UTF_8)); + } + } + } return setAuthorizationHeader(context, tokenRequestContext) .then(Mono.just(true)); @@ -214,38 +236,38 @@ public void authorizeRequestSync(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); setAuthorizationHeaderSync(context, tokenRequestContext); + return; } // 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. + // 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.getBodyAsBinaryData() != null) { context.setData(KEY_VAULT_STASHED_CONTENT_KEY, request.getBodyAsBinaryData()); context.setData(KEY_VAULT_STASHED_CONTENT_LENGTH_KEY, - request.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH)); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((BinaryData) null); } } } - @SuppressWarnings("unchecked") @Override public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) { HttpRequest request = context.getHttpRequest(); @@ -254,12 +276,12 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((BinaryData) (contentOptional.get())); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); Map challengeAttributes = - extractChallengeAttributes(response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + extractChallengeAttributes(response.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); String scope = challengeAttributes.get("resource"); if (scope != null) { @@ -309,12 +331,131 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims(new String(Base64Util.decodeString(claims))); + } + } + } setAuthorizationHeaderSync(context, tokenRequestContext); + return true; } + @Override + public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + return Mono.error(new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return authorizeRequest(context).then(Mono.defer(next::process)).flatMap(httpResponse -> { + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallenge(context, httpResponse, nextPolicy); + } + + return Mono.just(httpResponse); + }); + } + + @Override + public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + throw LOGGER.logExceptionAsError( + new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + + authorizeRequestSync(context); + + HttpResponse httpResponse = next.processSync(); + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallengeSync(context, httpResponse, nextPolicy); + } + + return httpResponse; + } + + private Mono handleChallenge(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextPolicy next) { + return authorizeRequestOnChallenge(context, httpResponse).flatMap(authorized -> { + if (authorized) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return next.process().flatMap(newResponse -> { + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallenge(context, newResponse, nextPolicy); + } else { + return Mono.just(newResponse); + } + }); + } + + return Mono.just(httpResponse); + }); + } + + private HttpResponse handleChallengeSync(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextSyncPolicy next) { + if (authorizeRequestOnChallengeSync(context, httpResponse)) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + HttpResponse newResponse = next.processSync(); + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallengeSync(context, newResponse, nextPolicy); + } + + return newResponse; + } + + return httpResponse; + } + + private boolean isClaimsPresent(HttpResponse httpResponse) { + Map challengeAttributes = + extractChallengeAttributes(httpResponse.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + + String error = challengeAttributes.get("error"); + + if (error != null) { + String base64Claims = challengeAttributes.get("claims"); + + return "insufficient_claims".equalsIgnoreCase(error) && base64Claims != null; + } + + return false; + } + private static class ChallengeParameters { private final URI authorizationUri; private final String tenantId; diff --git a/sdk/keyvault/azure-security-keyvault-administration/src/test/java/com/azure/security/keyvault/administration/KeyVaultCredentialPolicyTest.java b/sdk/keyvault/azure-security-keyvault-administration/src/test/java/com/azure/security/keyvault/administration/KeyVaultCredentialPolicyTest.java index 0e6eeb3876a78..b5caeb277a755 100644 --- a/sdk/keyvault/azure-security-keyvault-administration/src/test/java/com/azure/security/keyvault/administration/KeyVaultCredentialPolicyTest.java +++ b/sdk/keyvault/azure-security-keyvault-administration/src/test/java/com/azure/security/keyvault/administration/KeyVaultCredentialPolicyTest.java @@ -3,7 +3,9 @@ package com.azure.security.keyvault.administration; -import com.azure.core.credential.BasicAuthenticationCredential; +import com.azure.core.credential.AccessToken; +import com.azure.core.credential.TokenCredential; +import com.azure.core.credential.TokenRequestContext; import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; @@ -15,7 +17,9 @@ import com.azure.core.test.SyncAsyncExtension; import com.azure.core.test.annotation.SyncAsyncTest; import com.azure.core.test.http.MockHttpResponse; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; +import com.azure.core.util.Context; import com.azure.security.keyvault.administration.implementation.KeyVaultCredentialPolicy; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -28,12 +32,22 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; +import java.util.function.Function; import java.util.stream.Stream; +import static com.azure.core.http.HttpHeaderName.AUTHORIZATION; +import static com.azure.core.util.CoreUtils.isNullOrEmpty; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -43,85 +57,135 @@ public class KeyVaultCredentialPolicyTest { private static final String AUTHENTICATE_HEADER = "Bearer authorization=\"https://login.windows.net/72f988bf-86f1-41af-91ab-2d7cd022db57\", " + "resource=\"https://vault.azure.net\""; + private static final String AUTHENTICATE_HEADER_WITH_CLAIMS = + "Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", " + + "error=\"insufficient_claims\", " + + "claims=\"eyJhY2Nlc3NfdG9rZW4iOnsiYWNycyI6eyJlc3NlbnRpYWwiOnRydWUsInZhbHVlIjoiY3AxIn19fQ==\""; + private static final String DECODED_CLAIMS = "{\"access_token\":{\"acrs\":{\"essential\":true,\"value\":\"cp1\"}}}"; private static final String BEARER = "Bearer"; private static final String BODY = "this is a sample body"; private static final Flux BODY_FLUX = Flux.defer(() -> - Flux.fromStream( - Stream.of(BODY.split("")) - .map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8))) - )); - private BasicAuthenticationCredential credential; + Flux.fromStream(Stream.of(BODY.split("")).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8))))); + private static final String FAKE_ENCODED_CREDENTIAL = + Base64Util.encodeToString("user:fakePasswordPlaceholder".getBytes(StandardCharsets.UTF_8)); + private static final List> BASE_ASSERTIONS = Arrays.asList( + tokenRequestContext -> !tokenRequestContext.getScopes().isEmpty(), + tokenRequestContext -> !isNullOrEmpty(tokenRequestContext.getTenantId()), + TokenRequestContext::isCaeEnabled); + + private HttpResponse simpleResponse; + private HttpResponse unauthorizedHttpResponseWithWrongStatusCode; private HttpResponse unauthorizedHttpResponseWithHeader; private HttpResponse unauthorizedHttpResponseWithoutHeader; + private HttpResponse unauthorizedHttpResponseWithHeaderAndClaims; private HttpPipelineCallContext callContext; private HttpPipelineCallContext differentScopeContext; private HttpPipelineCallContext testContext; private HttpPipelineCallContext bodyContext; private HttpPipelineCallContext bodyFluxContext; + private TokenCredential credential; + + private static HttpPipelineCallContext createCallContext(HttpRequest request, Context context) { + AtomicReference callContextReference = new AtomicReference<>(); + + HttpPipeline callContextCreator = new HttpPipelineBuilder() + .policies((callContext, next) -> { + callContextReference.set(callContext); + + return next.process(); + }) + .httpClient(ignored -> Mono.empty()) + .build(); + + callContextCreator.sendSync(request, context); + + return callContextReference.get(); + } @BeforeEach public void setup() { HttpRequest request = new HttpRequest(HttpMethod.GET, "https://kvtest.vault.azure.net"); HttpRequest requestWithDifferentScope = new HttpRequest(HttpMethod.GET, "https://mytest.azurecr.io"); - HttpPipelineCallContext plainContext = createMockContext(request, null); + Context bodyContextContext = new Context("KeyVaultCredentialPolicyStashedBody", BinaryData.fromString(BODY)) + .addData("KeyVaultCredentialPolicyStashedContentLength", "21"); - HttpPipelineCallContext differentScopeContext = createMockContext(requestWithDifferentScope, null); + Context bodyFluxContextContext = new Context("KeyVaultCredentialPolicyStashedBody", BODY_FLUX) + .addData("KeyVaultCredentialPolicyStashedContentLength", "21"); - HttpPipelineCallContext testContext = createMockContext(request, null); + MockHttpResponse simpleResponse = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 200); - HttpPipelineCallContext bodyContext = createMockContext(request, callContext -> { - callContext.setData("KeyVaultCredentialPolicyStashedBody", BinaryData.fromString(BODY)); - callContext.setData("KeyVaultCredentialPolicyStashedContentLength", "21"); - }); + MockHttpResponse unauthorizedResponseWithWrongStatusCode = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); - HttpPipelineCallContext bodyFluxContext = createMockContext(request, callContext -> { - callContext.setData("KeyVaultCredentialPolicyStashedBody", BODY_FLUX); - callContext.setData("KeyVaultCredentialPolicyStashedContentLength", "21"); - }); + MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401); MockHttpResponse unauthorizedResponseWithHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://kvtest.vault.azure.net"), - 500, - new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER) - ); + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, + new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER)); - MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://kvtest.vault.azure.net"), 500); + MockHttpResponse unauthorizedResponseWithHeaderAndClaims = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, + new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER_WITH_CLAIMS)); + this.simpleResponse = simpleResponse; + this.unauthorizedHttpResponseWithWrongStatusCode = unauthorizedResponseWithWrongStatusCode; this.unauthorizedHttpResponseWithHeader = unauthorizedResponseWithHeader; this.unauthorizedHttpResponseWithoutHeader = unauthorizedResponseWithoutHeader; - this.callContext = plainContext; - this.differentScopeContext = differentScopeContext; - this.credential = new BasicAuthenticationCredential("user", "fakePasswordPlaceholder"); - this.testContext = testContext; - this.bodyContext = bodyContext; - this.bodyFluxContext = bodyFluxContext; + this.unauthorizedHttpResponseWithHeaderAndClaims = unauthorizedResponseWithHeaderAndClaims; + this.callContext = createCallContext(request, Context.NONE); + this.differentScopeContext = createCallContext(requestWithDifferentScope, Context.NONE); + this.testContext = createCallContext(request, Context.NONE); + this.bodyContext = createCallContext(request, bodyContextContext); + this.bodyFluxContext = createCallContext(request, bodyFluxContextContext); + // Can't use BasicAuthenticationCredential until the following PR is merged: + // https://github.com/Azure/azure-sdk-for-java/pull/42238 + this.credential = tokenRequestContext -> + Mono.fromCallable(() -> new AccessToken(FAKE_ENCODED_CREDENTIAL, OffsetDateTime.MAX.minusYears(1))); + } + @AfterEach + public void cleanup() { + KeyVaultCredentialPolicy.clearCache(); } - private static HttpPipelineCallContext createMockContext(HttpRequest request, - Consumer contextSetter) { - AtomicReference capturedContext = new AtomicReference<>(); - HttpPipeline httpPipeline = new HttpPipelineBuilder() - .policies((httpPipelineCallContext, httpPipelineNextPolicy) -> { - if (contextSetter != null) { - contextSetter.accept(httpPipelineCallContext); - } + @SyncAsyncTest + public void onNon401ErrorResponse() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); - capturedContext.set(httpPipelineCallContext); - return Mono.empty(); - }) - .httpClient(ignored -> Mono.empty()) + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithWrongStatusCode)) .build(); - httpPipeline.send(request).block(); + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); - return capturedContext.get(); + KeyVaultCredentialPolicy.clearCache(); } - @AfterEach - public void cleanup() { + @SyncAsyncTest + public void on401UnauthorizedResponseWithHeader() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithHeader)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNotNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + KeyVaultCredentialPolicy.clearCache(); } @@ -130,14 +194,14 @@ public void onChallengeCredentialPolicy() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); - String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } @@ -146,16 +210,125 @@ public void onChallengeCredentialPolicy() { public void onAuthorizeRequestChallengeCachePresent() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + StepVerifier.create(onChallengeAndClearCache(policy, this.callContext, unauthorizedHttpResponseWithHeader) // Challenge cache created + .then(policy.authorizeRequest(this.testContext))) // Challenge cache used + .verifyComplete(); + + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + assertFalse(tokenValue.isEmpty()); + assertTrue(tokenValue.startsWith(BEARER)); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Challenge cache created - onChallenge(policy, this.callContext, unauthorizedHttpResponseWithHeader).block(); + onChallengeAndClearCacheSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); // Challenge cache used - policy.authorizeRequest(this.testContext).block(); + policy.authorizeRequestSync(this.testContext); - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeader) + .flatMap(authorized -> { + if (authorized) { + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + return policy.authorizeRequestOnChallenge(this.callContext, // Challenge with claims received + this.unauthorizedHttpResponseWithHeaderAndClaims) + .map(ignored -> firstToken); + } else { + return Mono.just(""); + } + })) + .assertNext(firstToken -> { + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaims() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeaderAndClaims)) + .assertNext(result -> { + assertFalse(result); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaimsSync() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + // Challenge cache created + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, this.unauthorizedHttpResponseWithHeader)); + + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // Challenge with claims received + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaimsSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + // Challenge with claims received + assertFalse(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onAuthorizeRequestNoCache() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); @@ -166,7 +339,7 @@ public void onAuthorizeRequestNoCache() { () -> policy.authorizeRequest(this.callContext) ); - assertNull(this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION)); + assertNull(this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); } @SyncAsyncTest @@ -174,21 +347,21 @@ public void testSetContentLengthHeader() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); HttpHeaders headers = this.bodyFluxContext.getHttpRequest().getHeaders(); - String tokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = headers.getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); assertEquals("21", headers.getValue(HttpHeaderName.CONTENT_LENGTH)); HttpHeaders syncHeaders = this.bodyContext.getHttpRequest().getHeaders(); - String syncTokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String syncTokenValue = headers.getValue(AUTHORIZATION); assertFalse(syncTokenValue.isEmpty()); assertTrue(syncTokenValue.startsWith(BEARER)); assertEquals("21", syncHeaders.getValue(HttpHeaderName.CONTENT_LENGTH)); @@ -199,8 +372,8 @@ public void onAuthorizeRequestNoScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) ); assertFalse(onChallenge); @@ -211,9 +384,11 @@ public void onAuthorizeRequestDifferentScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); assertThrows(RuntimeException.class, - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)); + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)); - StepVerifier.create(onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)) + StepVerifier.create(onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)) .verifyErrorMessage("The challenge resource 'https://vault.azure.net/.default' does not match the " + "requested domain. If you wish to disable this check for your client, pass 'true' to the " + "SecretClientBuilder.disableChallengeResourceVerification() method when building it. See " @@ -225,38 +400,290 @@ public void onAuthorizeRequestDifferentScopeVerifyFalse() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, true); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader) ); assertTrue(onChallenge); } - @Test - public void onAuthorizeRequestChallengeCachePresentSync() { - KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Normal flow: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void processMultipleResponses() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); - // Challenge cache created - onChallengeSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); - // Challenge cache used - policy.authorizeRequestSync(this.testContext); + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); - assertFalse(tokenValue.isEmpty()); - assertTrue(tokenValue.startsWith(BEARER)); + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request. + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // On receiving an unauthorized response with claims, the token should be updated and a new attempt to make the + // original request should be made. + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the response with claims and set on the request. + assertNotNull(newToken); + // The token was updated. + assertNotEquals(firstToken, newToken); + // A subsequent request was successful. + assertEquals(simpleResponse, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized with claims + @SyncAsyncTest + public void processConsecutiveResponsesWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response with claims is received, it shall be returned as is. + unauthorizedHttpResponseWithHeaderAndClaims, + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeaderAndClaims, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized + @SyncAsyncTest + public void process401WithoutClaimsAfter401WithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response is received, it shall be returned as is. + unauthorizedHttpResponseWithHeader + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeader, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void process401WithClaimsAfter401WithoutClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)); + final String[] firstToken = new String[1]; + + testCredential.addAssertion(tokenRequestContext -> { + // This will ensure that that the first request does not contains claims, but the second does after + // receiving a 401 response with a challenge with claims. + testCredential.replaceAssertion(anotherTokenRequestContext -> + DECODED_CLAIMS.equals(anotherTokenRequestContext.getClaims()), 3); + + // We will also store the value of the first credential before it changes on a second call + firstToken[0] = Base64Util.encodeToString(testCredential.getCredential().getBytes(StandardCharsets.UTF_8)); + + assertNotNull(firstToken[0]); + + return tokenRequestContext.getClaims() == null; + }); + + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first unauthorized response caused a token to be set on the request, then the token was updated on a + // subsequent unauthorized response with claims. + assertNotEquals(firstToken[0], newToken); + // Finally, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + KeyVaultCredentialPolicy.clearCache(); } - private Mono onChallenge(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private Mono onChallengeAndClearCache(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { Mono onChallenge = policy.authorizeRequestOnChallenge(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallenge; } - private boolean onChallengeSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private boolean onChallengeAndClearCacheSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { boolean onChallengeSync = policy.authorizeRequestOnChallengeSync(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallengeSync; } + + private static final class MutableTestCredential implements TokenCredential { + private String credential; + private List> assertions; + + private MutableTestCredential(List> assertions) { + this.credential = new Random().toString(); + this.assertions = assertions; + } + + /** + * @throws RuntimeException if any of the assertions fail. + */ + @Override + public Mono getToken(TokenRequestContext requestContext) { + if (requestContext.isCaeEnabled() && requestContext.getClaims() != null) { + credential = new Random().toString(); + } + + String encodedCredential = Base64Util.encodeToString(credential.getBytes(StandardCharsets.UTF_8)); + + for (int i = 0; i < assertions.size(); i++) { + if (!assertions.get(i).apply(requestContext)) { + return Mono.error(new RuntimeException(String.format("Assertion number %d failed", i))); + } + } + + return Mono.fromCallable(() -> new AccessToken(encodedCredential, OffsetDateTime.MAX.minusYears(1))); + } + + private MutableTestCredential setAssertions(List> assertions) { + this.assertions = assertions; + + return this; + } + + private MutableTestCredential addAssertion(Function assertion) { + assertions.add(assertion); + + return this; + } + + private MutableTestCredential replaceAssertion(Function assertion, int index) { + assertions.set(index, assertion); + + return this; + } + + private String getCredential() { + return this.credential; + } + } } diff --git a/sdk/keyvault/azure-security-keyvault-certificates/CHANGELOG.md b/sdk/keyvault/azure-security-keyvault-certificates/CHANGELOG.md index dee1ba83b8456..bfad96806e473 100644 --- a/sdk/keyvault/azure-security-keyvault-certificates/CHANGELOG.md +++ b/sdk/keyvault/azure-security-keyvault-certificates/CHANGELOG.md @@ -2,6 +2,9 @@ ## 4.7.0 (2024-10-15) +## Features Added +- Added support for Continuous Access Evaluation (CAE). ([#41814](https://github.com/Azure/azure-sdk-for-java/pull/41814)) + ### Other Changes #### Dependency Updates diff --git a/sdk/keyvault/azure-security-keyvault-certificates/src/main/java/com/azure/security/keyvault/certificates/implementation/KeyVaultCredentialPolicy.java b/sdk/keyvault/azure-security-keyvault-certificates/src/main/java/com/azure/security/keyvault/certificates/implementation/KeyVaultCredentialPolicy.java index 6a2920aeddcfa..a978e2a585f8f 100644 --- a/sdk/keyvault/azure-security-keyvault-certificates/src/main/java/com/azure/security/keyvault/certificates/implementation/KeyVaultCredentialPolicy.java +++ b/sdk/keyvault/azure-security-keyvault-certificates/src/main/java/com/azure/security/keyvault/certificates/implementation/KeyVaultCredentialPolicy.java @@ -5,9 +5,12 @@ 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.HttpPipelineNextSyncPolicy; import com.azure.core.http.HttpRequest; import com.azure.core.http.HttpResponse; import com.azure.core.http.policy.BearerTokenAuthenticationPolicy; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; @@ -18,6 +21,7 @@ import java.net.URISyntaxException; import java.net.URL; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashMap; import java.util.Locale; @@ -26,6 +30,9 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import static com.azure.core.http.HttpHeaderName.CONTENT_LENGTH; +import static com.azure.core.http.HttpHeaderName.WWW_AUTHENTICATE; + /** * 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. @@ -35,10 +42,8 @@ public class KeyVaultCredentialPolicy extends BearerTokenAuthenticationPolicy { private static final ClientLogger LOGGER = new ClientLogger(KeyVaultCredentialPolicy.class); private static final String BEARER_TOKEN_PREFIX = "Bearer "; - 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 CHALLENGE_CACHE = new ConcurrentHashMap<>(); private ChallengeParameters challenge; private final boolean disableChallengeResourceVerification; @@ -68,16 +73,17 @@ private static Map extractChallengeAttributes(String authenticat return Collections.emptyMap(); } - authenticateHeader = - authenticateHeader.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), ""); - - String[] attributes = authenticateHeader.split(", "); + String[] attributes = authenticateHeader + .replace("\"", "") + .substring(authChallengePrefix.length()) + .split(","); Map attributeMap = new HashMap<>(); for (String pair : attributes) { - String[] keyValue = pair.split("="); + // Using trim is ugly, but we need it here because currently the 'claims' attribute comes after two spaces. + String[] keyValue = pair.trim().split("=", 2); - attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", "")); + attributeMap.put(keyValue[0], keyValue[1]); } return attributeMap; @@ -103,31 +109,31 @@ public Mono authorizeRequest(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); 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. + // 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.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((Flux) null); } } @@ -146,7 +152,7 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((Flux) contentOptional.get()); - request.setHeader(CONTENT_LENGTH_HEADER, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); @@ -201,7 +207,23 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims( + new String(Base64Util.decodeString(claims), StandardCharsets.UTF_8)); + } + } + } return setAuthorizationHeader(context, tokenRequestContext) .then(Mono.just(true)); @@ -214,38 +236,38 @@ public void authorizeRequestSync(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); setAuthorizationHeaderSync(context, tokenRequestContext); + return; } // 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. + // 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.getBodyAsBinaryData() != null) { context.setData(KEY_VAULT_STASHED_CONTENT_KEY, request.getBodyAsBinaryData()); context.setData(KEY_VAULT_STASHED_CONTENT_LENGTH_KEY, - request.getHeaders().getValue(CONTENT_LENGTH_HEADER)); - request.setHeader(CONTENT_LENGTH_HEADER, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((BinaryData) null); } } } - @SuppressWarnings("unchecked") @Override public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) { HttpRequest request = context.getHttpRequest(); @@ -254,7 +276,7 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((BinaryData) (contentOptional.get())); - request.setHeader(CONTENT_LENGTH_HEADER, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); @@ -309,12 +331,131 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims(new String(Base64Util.decodeString(claims))); + } + } + } setAuthorizationHeaderSync(context, tokenRequestContext); + return true; } + @Override + public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + return Mono.error(new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return authorizeRequest(context).then(Mono.defer(next::process)).flatMap(httpResponse -> { + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallenge(context, httpResponse, nextPolicy); + } + + return Mono.just(httpResponse); + }); + } + + @Override + public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + throw LOGGER.logExceptionAsError( + new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + + authorizeRequestSync(context); + + HttpResponse httpResponse = next.processSync(); + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallengeSync(context, httpResponse, nextPolicy); + } + + return httpResponse; + } + + private Mono handleChallenge(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextPolicy next) { + return authorizeRequestOnChallenge(context, httpResponse).flatMap(authorized -> { + if (authorized) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return next.process().flatMap(newResponse -> { + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallenge(context, newResponse, nextPolicy); + } else { + return Mono.just(newResponse); + } + }); + } + + return Mono.just(httpResponse); + }); + } + + private HttpResponse handleChallengeSync(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextSyncPolicy next) { + if (authorizeRequestOnChallengeSync(context, httpResponse)) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + HttpResponse newResponse = next.processSync(); + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallengeSync(context, newResponse, nextPolicy); + } + + return newResponse; + } + + return httpResponse; + } + + private boolean isClaimsPresent(HttpResponse httpResponse) { + Map challengeAttributes = + extractChallengeAttributes(httpResponse.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + + String error = challengeAttributes.get("error"); + + if (error != null) { + String base64Claims = challengeAttributes.get("claims"); + + return "insufficient_claims".equalsIgnoreCase(error) && base64Claims != null; + } + + return false; + } + private static class ChallengeParameters { private final URI authorizationUri; private final String tenantId; diff --git a/sdk/keyvault/azure-security-keyvault-certificates/src/test/java/com/azure/security/keyvault/certificates/KeyVaultCredentialPolicyTest.java b/sdk/keyvault/azure-security-keyvault-certificates/src/test/java/com/azure/security/keyvault/certificates/KeyVaultCredentialPolicyTest.java index 41a547a17179d..0dfbedd591fb5 100644 --- a/sdk/keyvault/azure-security-keyvault-certificates/src/test/java/com/azure/security/keyvault/certificates/KeyVaultCredentialPolicyTest.java +++ b/sdk/keyvault/azure-security-keyvault-certificates/src/test/java/com/azure/security/keyvault/certificates/KeyVaultCredentialPolicyTest.java @@ -3,7 +3,9 @@ package com.azure.security.keyvault.certificates; -import com.azure.core.credential.BasicAuthenticationCredential; +import com.azure.core.credential.AccessToken; +import com.azure.core.credential.TokenCredential; +import com.azure.core.credential.TokenRequestContext; import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; @@ -15,6 +17,7 @@ import com.azure.core.test.SyncAsyncExtension; import com.azure.core.test.annotation.SyncAsyncTest; import com.azure.core.test.http.MockHttpResponse; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.security.keyvault.certificates.implementation.KeyVaultCredentialPolicy; @@ -29,11 +32,22 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import java.util.stream.Stream; +import static com.azure.core.http.HttpHeaderName.AUTHORIZATION; +import static com.azure.core.util.CoreUtils.isNullOrEmpty; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -43,19 +57,33 @@ public class KeyVaultCredentialPolicyTest { private static final String AUTHENTICATE_HEADER = "Bearer authorization=\"https://login.windows.net/72f988bf-86f1-41af-91ab-2d7cd022db57\", " + "resource=\"https://vault.azure.net\""; - + private static final String AUTHENTICATE_HEADER_WITH_CLAIMS = + "Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", " + + "error=\"insufficient_claims\", " + + "claims=\"eyJhY2Nlc3NfdG9rZW4iOnsiYWNycyI6eyJlc3NlbnRpYWwiOnRydWUsInZhbHVlIjoiY3AxIn19fQ==\""; + private static final String DECODED_CLAIMS = "{\"access_token\":{\"acrs\":{\"essential\":true,\"value\":\"cp1\"}}}"; private static final String BEARER = "Bearer"; private static final String BODY = "this is a sample body"; private static final Flux BODY_FLUX = Flux.defer(() -> Flux.fromStream(Stream.of(BODY.split("")).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8))))); - private BasicAuthenticationCredential credential; + private static final String FAKE_ENCODED_CREDENTIAL = + Base64Util.encodeToString("user:fakePasswordPlaceholder".getBytes(StandardCharsets.UTF_8)); + private static final List> BASE_ASSERTIONS = Arrays.asList( + tokenRequestContext -> !tokenRequestContext.getScopes().isEmpty(), + tokenRequestContext -> !isNullOrEmpty(tokenRequestContext.getTenantId()), + TokenRequestContext::isCaeEnabled); + + private HttpResponse simpleResponse; + private HttpResponse unauthorizedHttpResponseWithWrongStatusCode; private HttpResponse unauthorizedHttpResponseWithHeader; private HttpResponse unauthorizedHttpResponseWithoutHeader; + private HttpResponse unauthorizedHttpResponseWithHeaderAndClaims; private HttpPipelineCallContext callContext; private HttpPipelineCallContext differentScopeContext; private HttpPipelineCallContext testContext; private HttpPipelineCallContext bodyContext; private HttpPipelineCallContext bodyFluxContext; + private TokenCredential credential; private static HttpPipelineCallContext createCallContext(HttpRequest request, Context context) { AtomicReference callContextReference = new AtomicReference<>(); @@ -63,12 +91,13 @@ private static HttpPipelineCallContext createCallContext(HttpRequest request, Co HttpPipeline callContextCreator = new HttpPipelineBuilder() .policies((callContext, next) -> { callContextReference.set(callContext); + return next.process(); }) .httpClient(ignored -> Mono.empty()) .build(); - callContextCreator.send(request, context).block(); + callContextCreator.sendSync(request, context); return callContextReference.get(); } @@ -84,22 +113,37 @@ public void setup() { Context bodyFluxContextContext = new Context("KeyVaultCredentialPolicyStashedBody", BODY_FLUX) .addData("KeyVaultCredentialPolicyStashedContentLength", "21"); + MockHttpResponse simpleResponse = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 200); + + MockHttpResponse unauthorizedResponseWithWrongStatusCode = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + + MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401); + MockHttpResponse unauthorizedResponseWithHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500, + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER)); - MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + MockHttpResponse unauthorizedResponseWithHeaderAndClaims = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, + new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER_WITH_CLAIMS)); + this.simpleResponse = simpleResponse; + this.unauthorizedHttpResponseWithWrongStatusCode = unauthorizedResponseWithWrongStatusCode; this.unauthorizedHttpResponseWithHeader = unauthorizedResponseWithHeader; this.unauthorizedHttpResponseWithoutHeader = unauthorizedResponseWithoutHeader; + this.unauthorizedHttpResponseWithHeaderAndClaims = unauthorizedResponseWithHeaderAndClaims; this.callContext = createCallContext(request, Context.NONE); this.differentScopeContext = createCallContext(requestWithDifferentScope, Context.NONE); - this.credential = new BasicAuthenticationCredential("user", "fakePasswordPlaceholder"); this.testContext = createCallContext(request, Context.NONE); this.bodyContext = createCallContext(request, bodyContextContext); this.bodyFluxContext = createCallContext(request, bodyFluxContextContext); - + // Can't use BasicAuthenticationCredential until the following PR is merged: + // https://github.com/Azure/azure-sdk-for-java/pull/42238 + this.credential = tokenRequestContext -> + Mono.fromCallable(() -> new AccessToken(FAKE_ENCODED_CREDENTIAL, OffsetDateTime.MAX.minusYears(1))); } @AfterEach @@ -107,19 +151,57 @@ public void cleanup() { KeyVaultCredentialPolicy.clearCache(); } + @SyncAsyncTest + public void onNon401ErrorResponse() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithWrongStatusCode)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + + @SyncAsyncTest + public void on401UnauthorizedResponseWithHeader() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithHeader)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNotNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onChallengeCredentialPolicy() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); - String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } @@ -128,16 +210,125 @@ public void onChallengeCredentialPolicy() { public void onAuthorizeRequestChallengeCachePresent() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + StepVerifier.create(onChallengeAndClearCache(policy, this.callContext, unauthorizedHttpResponseWithHeader) // Challenge cache created + .then(policy.authorizeRequest(this.testContext))) // Challenge cache used + .verifyComplete(); + + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + assertFalse(tokenValue.isEmpty()); + assertTrue(tokenValue.startsWith(BEARER)); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Challenge cache created - onChallenge(policy, this.callContext, unauthorizedHttpResponseWithHeader).block(); + onChallengeAndClearCacheSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); // Challenge cache used - policy.authorizeRequest(this.testContext).block(); + policy.authorizeRequestSync(this.testContext); - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeader) + .flatMap(authorized -> { + if (authorized) { + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + return policy.authorizeRequestOnChallenge(this.callContext, // Challenge with claims received + this.unauthorizedHttpResponseWithHeaderAndClaims) + .map(ignored -> firstToken); + } else { + return Mono.just(""); + } + })) + .assertNext(firstToken -> { + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaims() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeaderAndClaims)) + .assertNext(result -> { + assertFalse(result); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaimsSync() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + // Challenge cache created + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, this.unauthorizedHttpResponseWithHeader)); + + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // Challenge with claims received + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaimsSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + // Challenge with claims received + assertFalse(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onAuthorizeRequestNoCache() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); @@ -148,7 +339,7 @@ public void onAuthorizeRequestNoCache() { () -> policy.authorizeRequest(this.callContext) ); - assertNull(this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION)); + assertNull(this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); } @SyncAsyncTest @@ -156,21 +347,21 @@ public void testSetContentLengthHeader() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); HttpHeaders headers = this.bodyFluxContext.getHttpRequest().getHeaders(); - String tokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = headers.getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); assertEquals("21", headers.getValue(HttpHeaderName.CONTENT_LENGTH)); HttpHeaders syncHeaders = this.bodyContext.getHttpRequest().getHeaders(); - String syncTokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String syncTokenValue = headers.getValue(AUTHORIZATION); assertFalse(syncTokenValue.isEmpty()); assertTrue(syncTokenValue.startsWith(BEARER)); assertEquals("21", syncHeaders.getValue(HttpHeaderName.CONTENT_LENGTH)); @@ -181,8 +372,8 @@ public void onAuthorizeRequestNoScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) ); assertFalse(onChallenge); @@ -193,9 +384,11 @@ public void onAuthorizeRequestDifferentScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); assertThrows(RuntimeException.class, - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)); + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)); - StepVerifier.create(onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)) + StepVerifier.create(onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)) .verifyErrorMessage("The challenge resource 'https://vault.azure.net/.default' does not match the " + "requested domain. If you wish to disable this check for your client, pass 'true' to the " + "SecretClientBuilder.disableChallengeResourceVerification() method when building it. See " @@ -207,38 +400,290 @@ public void onAuthorizeRequestDifferentScopeVerifyFalse() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, true); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader) ); assertTrue(onChallenge); } - @Test - public void onAuthorizeRequestChallengeCachePresentSync() { - KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Normal flow: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void processMultipleResponses() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); - // Challenge cache created - onChallengeSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); - // Challenge cache used - policy.authorizeRequestSync(this.testContext); + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); - assertFalse(tokenValue.isEmpty()); - assertTrue(tokenValue.startsWith(BEARER)); + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request. + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // On receiving an unauthorized response with claims, the token should be updated and a new attempt to make the + // original request should be made. + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the response with claims and set on the request. + assertNotNull(newToken); + // The token was updated. + assertNotEquals(firstToken, newToken); + // A subsequent request was successful. + assertEquals(simpleResponse, newResponse); + + KeyVaultCredentialPolicy.clearCache(); } - private Mono onChallenge(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized with claims + @SyncAsyncTest + public void processConsecutiveResponsesWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response with claims is received, it shall be returned as is. + unauthorizedHttpResponseWithHeaderAndClaims, + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeaderAndClaims, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized + @SyncAsyncTest + public void process401WithoutClaimsAfter401WithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response is received, it shall be returned as is. + unauthorizedHttpResponseWithHeader + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeader, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void process401WithClaimsAfter401WithoutClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)); + final String[] firstToken = new String[1]; + + testCredential.addAssertion(tokenRequestContext -> { + // This will ensure that that the first request does not contains claims, but the second does after + // receiving a 401 response with a challenge with claims. + testCredential.replaceAssertion(anotherTokenRequestContext -> + DECODED_CLAIMS.equals(anotherTokenRequestContext.getClaims()), 3); + + // We will also store the value of the first credential before it changes on a second call + firstToken[0] = Base64Util.encodeToString(testCredential.getCredential().getBytes(StandardCharsets.UTF_8)); + + assertNotNull(firstToken[0]); + + return tokenRequestContext.getClaims() == null; + }); + + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first unauthorized response caused a token to be set on the request, then the token was updated on a + // subsequent unauthorized response with claims. + assertNotEquals(firstToken[0], newToken); + // Finally, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + private Mono onChallengeAndClearCache(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { Mono onChallenge = policy.authorizeRequestOnChallenge(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallenge; } - private boolean onChallengeSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private boolean onChallengeAndClearCacheSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { boolean onChallengeSync = policy.authorizeRequestOnChallengeSync(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallengeSync; } + + private static final class MutableTestCredential implements TokenCredential { + private String credential; + private List> assertions; + + private MutableTestCredential(List> assertions) { + this.credential = new Random().toString(); + this.assertions = assertions; + } + + /** + * @throws RuntimeException if any of the assertions fail. + */ + @Override + public Mono getToken(TokenRequestContext requestContext) { + if (requestContext.isCaeEnabled() && requestContext.getClaims() != null) { + credential = new Random().toString(); + } + + String encodedCredential = Base64Util.encodeToString(credential.getBytes(StandardCharsets.UTF_8)); + + for (int i = 0; i < assertions.size(); i++) { + if (!assertions.get(i).apply(requestContext)) { + return Mono.error(new RuntimeException(String.format("Assertion number %d failed", i))); + } + } + + return Mono.fromCallable(() -> new AccessToken(encodedCredential, OffsetDateTime.MAX.minusYears(1))); + } + + private MutableTestCredential setAssertions(List> assertions) { + this.assertions = assertions; + + return this; + } + + private MutableTestCredential addAssertion(Function assertion) { + assertions.add(assertion); + + return this; + } + + private MutableTestCredential replaceAssertion(Function assertion, int index) { + assertions.set(index, assertion); + + return this; + } + + private String getCredential() { + return this.credential; + } + } } diff --git a/sdk/keyvault/azure-security-keyvault-keys/CHANGELOG.md b/sdk/keyvault/azure-security-keyvault-keys/CHANGELOG.md index 39079781ed019..a0da4d28c739e 100644 --- a/sdk/keyvault/azure-security-keyvault-keys/CHANGELOG.md +++ b/sdk/keyvault/azure-security-keyvault-keys/CHANGELOG.md @@ -2,6 +2,7 @@ ## 4.9.0 (2024-10-15) - Added a new configuration flag to cryptography clients that allows deferring all cryptographic operations to the Key Vault service. ([#40384](https://github.com/Azure/azure-sdk-for-java/pull/40384)) +- Added support for Continuous Access Evaluation (CAE). ([#41814](https://github.com/Azure/azure-sdk-for-java/pull/41814)) ### Other Changes @@ -39,6 +40,19 @@ - Upgraded `azure-json` from `1.1.0` to version `1.2.0`. - Upgraded `azure-core` from `1.49.1` to version `1.50.0`. +## 4.9.0-beta.1 (2024-07-29) + +### Features Added +- Added a new configuration flag to cryptography clients to defer all cryptographic operations to the Key Vault service. ([#40384](https://github.com/Azure/azure-sdk-for-java/pull/40384)) + +### Other Changes + +#### Dependency Updates + +- Upgraded `azure-core-http-netty` from `1.15.1` to version `1.15.2`. +- Upgraded `azure-json` from `1.1.0` to version `1.2.0`. +- Upgraded `azure-core` from `1.49.1` to version `1.50.0`. + ## 4.8.5 (2024-06-27) ### Other Changes @@ -778,7 +792,6 @@ Changes when compared to the last stable release (`4.7.3`) include: - `KeyEncryptionKeyClientBuilder.buildKeyEncryptionKey` and `KeyEncryptionKeyClientBuilder.buildAsyncKeyEncryptionKey`supports consumption of a secret id representing the symmetric key stored in the Key Vault as a secret. - Dropped third party dependency on apache commons codec library. - ### Breaking changes - Key has been renamed to KeyVaultKey to avoid ambiguity with other libraries and to yield better search results. - Key.keyMaterial has been renamed to KeyVaultKey.key. diff --git a/sdk/keyvault/azure-security-keyvault-keys/src/main/java/com/azure/security/keyvault/keys/implementation/KeyVaultCredentialPolicy.java b/sdk/keyvault/azure-security-keyvault-keys/src/main/java/com/azure/security/keyvault/keys/implementation/KeyVaultCredentialPolicy.java index 8bfc86fafd5e3..608a4521d2459 100644 --- a/sdk/keyvault/azure-security-keyvault-keys/src/main/java/com/azure/security/keyvault/keys/implementation/KeyVaultCredentialPolicy.java +++ b/sdk/keyvault/azure-security-keyvault-keys/src/main/java/com/azure/security/keyvault/keys/implementation/KeyVaultCredentialPolicy.java @@ -4,11 +4,13 @@ import com.azure.core.credential.TokenCredential; import com.azure.core.credential.TokenRequestContext; -import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpPipelineCallContext; +import com.azure.core.http.HttpPipelineNextPolicy; +import com.azure.core.http.HttpPipelineNextSyncPolicy; import com.azure.core.http.HttpRequest; import com.azure.core.http.HttpResponse; import com.azure.core.http.policy.BearerTokenAuthenticationPolicy; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; @@ -19,6 +21,7 @@ import java.net.URISyntaxException; import java.net.URL; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashMap; import java.util.Locale; @@ -27,6 +30,9 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import static com.azure.core.http.HttpHeaderName.CONTENT_LENGTH; +import static com.azure.core.http.HttpHeaderName.WWW_AUTHENTICATE; + /** * 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. @@ -67,16 +73,17 @@ private static Map extractChallengeAttributes(String authenticat return Collections.emptyMap(); } - authenticateHeader = - authenticateHeader.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), ""); - - String[] attributes = authenticateHeader.split(", "); + String[] attributes = authenticateHeader + .replace("\"", "") + .substring(authChallengePrefix.length()) + .split(","); Map attributeMap = new HashMap<>(); for (String pair : attributes) { - String[] keyValue = pair.split("="); + // Using trim is ugly, but we need it here because currently the 'claims' attribute comes after two spaces. + String[] keyValue = pair.trim().split("=", 2); - attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", "")); + attributeMap.put(keyValue[0], keyValue[1]); } return attributeMap; @@ -102,31 +109,31 @@ public Mono authorizeRequest(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); 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. + // 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(HttpHeaderName.CONTENT_LENGTH)); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((Flux) null); } } @@ -145,12 +152,12 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((Flux) contentOptional.get()); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); - Map challengeAttributes = extractChallengeAttributes( - response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + Map challengeAttributes = + extractChallengeAttributes(response.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); String scope = challengeAttributes.get("resource"); if (scope != null) { @@ -200,7 +207,23 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims( + new String(Base64Util.decodeString(claims), StandardCharsets.UTF_8)); + } + } + } return setAuthorizationHeader(context, tokenRequestContext) .then(Mono.just(true)); @@ -213,38 +236,38 @@ public void authorizeRequestSync(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); setAuthorizationHeaderSync(context, tokenRequestContext); + return; } // 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. + // 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.getBodyAsBinaryData() != null) { context.setData(KEY_VAULT_STASHED_CONTENT_KEY, request.getBodyAsBinaryData()); context.setData(KEY_VAULT_STASHED_CONTENT_LENGTH_KEY, - request.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH)); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((BinaryData) null); } } } - @SuppressWarnings("unchecked") @Override public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) { HttpRequest request = context.getHttpRequest(); @@ -253,12 +276,12 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((BinaryData) (contentOptional.get())); - request.setHeader(HttpHeaderName.CONTENT_LENGTH, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); Map challengeAttributes = - extractChallengeAttributes(response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + extractChallengeAttributes(response.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); String scope = challengeAttributes.get("resource"); if (scope != null) { @@ -308,12 +331,131 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims(new String(Base64Util.decodeString(claims))); + } + } + } setAuthorizationHeaderSync(context, tokenRequestContext); + return true; } + @Override + public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + return Mono.error(new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return authorizeRequest(context).then(Mono.defer(next::process)).flatMap(httpResponse -> { + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallenge(context, httpResponse, nextPolicy); + } + + return Mono.just(httpResponse); + }); + } + + @Override + public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + throw LOGGER.logExceptionAsError( + new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + + authorizeRequestSync(context); + + HttpResponse httpResponse = next.processSync(); + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallengeSync(context, httpResponse, nextPolicy); + } + + return httpResponse; + } + + private Mono handleChallenge(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextPolicy next) { + return authorizeRequestOnChallenge(context, httpResponse).flatMap(authorized -> { + if (authorized) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return next.process().flatMap(newResponse -> { + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallenge(context, newResponse, nextPolicy); + } else { + return Mono.just(newResponse); + } + }); + } + + return Mono.just(httpResponse); + }); + } + + private HttpResponse handleChallengeSync(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextSyncPolicy next) { + if (authorizeRequestOnChallengeSync(context, httpResponse)) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + HttpResponse newResponse = next.processSync(); + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallengeSync(context, newResponse, nextPolicy); + } + + return newResponse; + } + + return httpResponse; + } + + private boolean isClaimsPresent(HttpResponse httpResponse) { + Map challengeAttributes = + extractChallengeAttributes(httpResponse.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + + String error = challengeAttributes.get("error"); + + if (error != null) { + String base64Claims = challengeAttributes.get("claims"); + + return "insufficient_claims".equalsIgnoreCase(error) && base64Claims != null; + } + + return false; + } + private static class ChallengeParameters { private final URI authorizationUri; private final String tenantId; diff --git a/sdk/keyvault/azure-security-keyvault-keys/src/test/java/com/azure/security/keyvault/keys/KeyVaultCredentialPolicyTest.java b/sdk/keyvault/azure-security-keyvault-keys/src/test/java/com/azure/security/keyvault/keys/KeyVaultCredentialPolicyTest.java index f63e9eb841ed1..b7fdc8d2600a9 100644 --- a/sdk/keyvault/azure-security-keyvault-keys/src/test/java/com/azure/security/keyvault/keys/KeyVaultCredentialPolicyTest.java +++ b/sdk/keyvault/azure-security-keyvault-keys/src/test/java/com/azure/security/keyvault/keys/KeyVaultCredentialPolicyTest.java @@ -3,7 +3,9 @@ package com.azure.security.keyvault.keys; -import com.azure.core.credential.BasicAuthenticationCredential; +import com.azure.core.credential.AccessToken; +import com.azure.core.credential.TokenCredential; +import com.azure.core.credential.TokenRequestContext; import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; @@ -15,6 +17,7 @@ import com.azure.core.test.SyncAsyncExtension; import com.azure.core.test.annotation.SyncAsyncTest; import com.azure.core.test.http.MockHttpResponse; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.security.keyvault.keys.implementation.KeyVaultCredentialPolicy; @@ -29,11 +32,22 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import java.util.stream.Stream; +import static com.azure.core.http.HttpHeaderName.AUTHORIZATION; +import static com.azure.core.util.CoreUtils.isNullOrEmpty; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -43,18 +57,33 @@ public class KeyVaultCredentialPolicyTest { private static final String AUTHENTICATE_HEADER = "Bearer authorization=\"https://login.windows.net/72f988bf-86f1-41af-91ab-2d7cd022db57\", " + "resource=\"https://vault.azure.net\""; + private static final String AUTHENTICATE_HEADER_WITH_CLAIMS = + "Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", " + + "error=\"insufficient_claims\", " + + "claims=\"eyJhY2Nlc3NfdG9rZW4iOnsiYWNycyI6eyJlc3NlbnRpYWwiOnRydWUsInZhbHVlIjoiY3AxIn19fQ==\""; + private static final String DECODED_CLAIMS = "{\"access_token\":{\"acrs\":{\"essential\":true,\"value\":\"cp1\"}}}"; private static final String BEARER = "Bearer"; private static final String BODY = "this is a sample body"; private static final Flux BODY_FLUX = Flux.defer(() -> Flux.fromStream(Stream.of(BODY.split("")).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8))))); - private BasicAuthenticationCredential credential; + private static final String FAKE_ENCODED_CREDENTIAL = + Base64Util.encodeToString("user:fakePasswordPlaceholder".getBytes(StandardCharsets.UTF_8)); + private static final List> BASE_ASSERTIONS = Arrays.asList( + tokenRequestContext -> !tokenRequestContext.getScopes().isEmpty(), + tokenRequestContext -> !isNullOrEmpty(tokenRequestContext.getTenantId()), + TokenRequestContext::isCaeEnabled); + + private HttpResponse simpleResponse; + private HttpResponse unauthorizedHttpResponseWithWrongStatusCode; private HttpResponse unauthorizedHttpResponseWithHeader; private HttpResponse unauthorizedHttpResponseWithoutHeader; + private HttpResponse unauthorizedHttpResponseWithHeaderAndClaims; private HttpPipelineCallContext callContext; private HttpPipelineCallContext differentScopeContext; private HttpPipelineCallContext testContext; private HttpPipelineCallContext bodyContext; private HttpPipelineCallContext bodyFluxContext; + private TokenCredential credential; private static HttpPipelineCallContext createCallContext(HttpRequest request, Context context) { AtomicReference callContextReference = new AtomicReference<>(); @@ -62,12 +91,13 @@ private static HttpPipelineCallContext createCallContext(HttpRequest request, Co HttpPipeline callContextCreator = new HttpPipelineBuilder() .policies((callContext, next) -> { callContextReference.set(callContext); + return next.process(); }) .httpClient(ignored -> Mono.empty()) .build(); - callContextCreator.send(request, context).block(); + callContextCreator.sendSync(request, context); return callContextReference.get(); } @@ -83,22 +113,37 @@ public void setup() { Context bodyFluxContextContext = new Context("KeyVaultCredentialPolicyStashedBody", BODY_FLUX) .addData("KeyVaultCredentialPolicyStashedContentLength", "21"); + MockHttpResponse simpleResponse = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 200); + + MockHttpResponse unauthorizedResponseWithWrongStatusCode = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + + MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401); + MockHttpResponse unauthorizedResponseWithHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500, + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER)); - MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + MockHttpResponse unauthorizedResponseWithHeaderAndClaims = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, + new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER_WITH_CLAIMS)); + this.simpleResponse = simpleResponse; + this.unauthorizedHttpResponseWithWrongStatusCode = unauthorizedResponseWithWrongStatusCode; this.unauthorizedHttpResponseWithHeader = unauthorizedResponseWithHeader; this.unauthorizedHttpResponseWithoutHeader = unauthorizedResponseWithoutHeader; + this.unauthorizedHttpResponseWithHeaderAndClaims = unauthorizedResponseWithHeaderAndClaims; this.callContext = createCallContext(request, Context.NONE); this.differentScopeContext = createCallContext(requestWithDifferentScope, Context.NONE); - this.credential = new BasicAuthenticationCredential("user", "fakePasswordPlaceholder"); this.testContext = createCallContext(request, Context.NONE); this.bodyContext = createCallContext(request, bodyContextContext); this.bodyFluxContext = createCallContext(request, bodyFluxContextContext); - + // Can't use BasicAuthenticationCredential until the following PR is merged: + // https://github.com/Azure/azure-sdk-for-java/pull/42238 + this.credential = tokenRequestContext -> + Mono.fromCallable(() -> new AccessToken(FAKE_ENCODED_CREDENTIAL, OffsetDateTime.MAX.minusYears(1))); } @AfterEach @@ -106,19 +151,57 @@ public void cleanup() { KeyVaultCredentialPolicy.clearCache(); } + @SyncAsyncTest + public void onNon401ErrorResponse() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithWrongStatusCode)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + + @SyncAsyncTest + public void on401UnauthorizedResponseWithHeader() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithHeader)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNotNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onChallengeCredentialPolicy() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); - String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } @@ -127,16 +210,125 @@ public void onChallengeCredentialPolicy() { public void onAuthorizeRequestChallengeCachePresent() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + StepVerifier.create(onChallengeAndClearCache(policy, this.callContext, unauthorizedHttpResponseWithHeader) // Challenge cache created + .then(policy.authorizeRequest(this.testContext))) // Challenge cache used + .verifyComplete(); + + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + assertFalse(tokenValue.isEmpty()); + assertTrue(tokenValue.startsWith(BEARER)); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Challenge cache created - onChallenge(policy, this.callContext, unauthorizedHttpResponseWithHeader).block(); + onChallengeAndClearCacheSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); // Challenge cache used - policy.authorizeRequest(this.testContext).block(); + policy.authorizeRequestSync(this.testContext); - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeader) + .flatMap(authorized -> { + if (authorized) { + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + return policy.authorizeRequestOnChallenge(this.callContext, // Challenge with claims received + this.unauthorizedHttpResponseWithHeaderAndClaims) + .map(ignored -> firstToken); + } else { + return Mono.just(""); + } + })) + .assertNext(firstToken -> { + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaims() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeaderAndClaims)) + .assertNext(result -> { + assertFalse(result); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaimsSync() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + // Challenge cache created + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, this.unauthorizedHttpResponseWithHeader)); + + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // Challenge with claims received + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaimsSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + // Challenge with claims received + assertFalse(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onAuthorizeRequestNoCache() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); @@ -147,7 +339,7 @@ public void onAuthorizeRequestNoCache() { () -> policy.authorizeRequest(this.callContext) ); - assertNull(this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION)); + assertNull(this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); } @SyncAsyncTest @@ -155,21 +347,21 @@ public void testSetContentLengthHeader() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); HttpHeaders headers = this.bodyFluxContext.getHttpRequest().getHeaders(); - String tokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = headers.getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); assertEquals("21", headers.getValue(HttpHeaderName.CONTENT_LENGTH)); HttpHeaders syncHeaders = this.bodyContext.getHttpRequest().getHeaders(); - String syncTokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String syncTokenValue = headers.getValue(AUTHORIZATION); assertFalse(syncTokenValue.isEmpty()); assertTrue(syncTokenValue.startsWith(BEARER)); assertEquals("21", syncHeaders.getValue(HttpHeaderName.CONTENT_LENGTH)); @@ -180,8 +372,8 @@ public void onAuthorizeRequestNoScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) ); assertFalse(onChallenge); @@ -192,9 +384,11 @@ public void onAuthorizeRequestDifferentScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); assertThrows(RuntimeException.class, - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)); + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)); - StepVerifier.create(onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)) + StepVerifier.create(onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)) .verifyErrorMessage("The challenge resource 'https://vault.azure.net/.default' does not match the " + "requested domain. If you wish to disable this check for your client, pass 'true' to the " + "SecretClientBuilder.disableChallengeResourceVerification() method when building it. See " @@ -206,38 +400,290 @@ public void onAuthorizeRequestDifferentScopeVerifyFalse() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, true); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader) ); assertTrue(onChallenge); } - @Test - public void onAuthorizeRequestChallengeCachePresentSync() { - KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Normal flow: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void processMultipleResponses() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); - // Challenge cache created - onChallengeSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); - // Challenge cache used - policy.authorizeRequestSync(this.testContext); + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); - assertFalse(tokenValue.isEmpty()); - assertTrue(tokenValue.startsWith(BEARER)); + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request. + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // On receiving an unauthorized response with claims, the token should be updated and a new attempt to make the + // original request should be made. + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the response with claims and set on the request. + assertNotNull(newToken); + // The token was updated. + assertNotEquals(firstToken, newToken); + // A subsequent request was successful. + assertEquals(simpleResponse, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized with claims + @SyncAsyncTest + public void processConsecutiveResponsesWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response with claims is received, it shall be returned as is. + unauthorizedHttpResponseWithHeaderAndClaims, + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeaderAndClaims, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized + @SyncAsyncTest + public void process401WithoutClaimsAfter401WithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response is received, it shall be returned as is. + unauthorizedHttpResponseWithHeader + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeader, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void process401WithClaimsAfter401WithoutClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)); + final String[] firstToken = new String[1]; + + testCredential.addAssertion(tokenRequestContext -> { + // This will ensure that that the first request does not contains claims, but the second does after + // receiving a 401 response with a challenge with claims. + testCredential.replaceAssertion(anotherTokenRequestContext -> + DECODED_CLAIMS.equals(anotherTokenRequestContext.getClaims()), 3); + + // We will also store the value of the first credential before it changes on a second call + firstToken[0] = Base64Util.encodeToString(testCredential.getCredential().getBytes(StandardCharsets.UTF_8)); + + assertNotNull(firstToken[0]); + + return tokenRequestContext.getClaims() == null; + }); + + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first unauthorized response caused a token to be set on the request, then the token was updated on a + // subsequent unauthorized response with claims. + assertNotEquals(firstToken[0], newToken); + // Finally, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + KeyVaultCredentialPolicy.clearCache(); } - private Mono onChallenge(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private Mono onChallengeAndClearCache(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { Mono onChallenge = policy.authorizeRequestOnChallenge(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallenge; } - private boolean onChallengeSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private boolean onChallengeAndClearCacheSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { boolean onChallengeSync = policy.authorizeRequestOnChallengeSync(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallengeSync; } + + private static final class MutableTestCredential implements TokenCredential { + private String credential; + private List> assertions; + + private MutableTestCredential(List> assertions) { + this.credential = new Random().toString(); + this.assertions = assertions; + } + + /** + * @throws RuntimeException if any of the assertions fail. + */ + @Override + public Mono getToken(TokenRequestContext requestContext) { + if (requestContext.isCaeEnabled() && requestContext.getClaims() != null) { + credential = new Random().toString(); + } + + String encodedCredential = Base64Util.encodeToString(credential.getBytes(StandardCharsets.UTF_8)); + + for (int i = 0; i < assertions.size(); i++) { + if (!assertions.get(i).apply(requestContext)) { + return Mono.error(new RuntimeException(String.format("Assertion number %d failed", i))); + } + } + + return Mono.fromCallable(() -> new AccessToken(encodedCredential, OffsetDateTime.MAX.minusYears(1))); + } + + private MutableTestCredential setAssertions(List> assertions) { + this.assertions = assertions; + + return this; + } + + private MutableTestCredential addAssertion(Function assertion) { + assertions.add(assertion); + + return this; + } + + private MutableTestCredential replaceAssertion(Function assertion, int index) { + assertions.set(index, assertion); + + return this; + } + + private String getCredential() { + return this.credential; + } + } } diff --git a/sdk/keyvault/azure-security-keyvault-secrets/CHANGELOG.md b/sdk/keyvault/azure-security-keyvault-secrets/CHANGELOG.md index 412efdf2d6225..37a365479aa27 100644 --- a/sdk/keyvault/azure-security-keyvault-secrets/CHANGELOG.md +++ b/sdk/keyvault/azure-security-keyvault-secrets/CHANGELOG.md @@ -19,7 +19,6 @@ - Upgraded `azure-core-http-netty` from `1.15.3` to version `1.15.4`. - Upgraded `azure-core` from `1.51.0` to version `1.52.0`. - ## 4.8.6 (2024-08-24) ### Other Changes @@ -29,7 +28,6 @@ - Upgraded `azure-core` from `1.50.0` to version `1.51.0`. - Upgraded `azure-core-http-netty` from `1.15.2` to version `1.15.3`. - ## 4.8.5 (2024-07-29) ### Other Changes @@ -40,7 +38,6 @@ - Upgraded `azure-json` from `1.1.0` to version `1.2.0`. - Upgraded `azure-core` from `1.49.1` to version `1.50.0`. - ## 4.8.4 (2024-06-27) ### Other Changes @@ -50,7 +47,6 @@ - Upgraded `azure-core` from `1.49.0` to version `1.49.1`. - Upgraded `azure-core-http-netty` from `1.15.0` to version `1.15.1`. - ## 4.8.3 (2024-05-13) ### Other Changes @@ -69,7 +65,6 @@ - Upgraded `azure-core` from `1.47.0` to version `1.48.0`. - Upgraded `azure-core-http-netty` from `1.14.1` to version `1.14.2`. - ## 4.8.1 (2024-03-20) ### Other Changes @@ -79,7 +74,6 @@ - Upgraded `azure-core` from `1.46.0` to version `1.47.0`. - Upgraded `azure-core-http-netty` from `1.14.0` to version `1.14.1`. - ## 4.8.0 (2024-02-22) Changes when compared to the last stable release (`4.7.3`) include: @@ -414,7 +408,6 @@ Changes when compared to the last stable release (`4.7.3`) include: - Upgraded `azure-core` dependency to `1.19.0` - Upgraded `azure-core-http-netty` dependency to `1.10.2` - ## 4.3.1 (2021-07-08) ### Other Changes @@ -630,7 +623,6 @@ and [samples](https://github.com/Azure/azure-sdk-for-java/blob/azure-keyvault-secrets_4.0.0-preview.1/keyvault/client/secrets/src/samples/java) demonstrate the new API. - ### Major changes from `azure-keyvault` - Packages scoped by functionality - `azure-keyvault-secrets` contains a `SecretClient` and `SecretAsyncClient` for secret operations, diff --git a/sdk/keyvault/azure-security-keyvault-secrets/src/main/java/com/azure/security/keyvault/secrets/implementation/KeyVaultCredentialPolicy.java b/sdk/keyvault/azure-security-keyvault-secrets/src/main/java/com/azure/security/keyvault/secrets/implementation/KeyVaultCredentialPolicy.java index 536a09fc114da..804564bc040e0 100644 --- a/sdk/keyvault/azure-security-keyvault-secrets/src/main/java/com/azure/security/keyvault/secrets/implementation/KeyVaultCredentialPolicy.java +++ b/sdk/keyvault/azure-security-keyvault-secrets/src/main/java/com/azure/security/keyvault/secrets/implementation/KeyVaultCredentialPolicy.java @@ -5,9 +5,12 @@ 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.HttpPipelineNextSyncPolicy; import com.azure.core.http.HttpRequest; import com.azure.core.http.HttpResponse; import com.azure.core.http.policy.BearerTokenAuthenticationPolicy; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; @@ -18,6 +21,7 @@ import java.net.URISyntaxException; import java.net.URL; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashMap; import java.util.Locale; @@ -26,6 +30,9 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import static com.azure.core.http.HttpHeaderName.CONTENT_LENGTH; +import static com.azure.core.http.HttpHeaderName.WWW_AUTHENTICATE; + /** * 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. @@ -35,10 +42,8 @@ public class KeyVaultCredentialPolicy extends BearerTokenAuthenticationPolicy { private static final ClientLogger LOGGER = new ClientLogger(KeyVaultCredentialPolicy.class); private static final String BEARER_TOKEN_PREFIX = "Bearer "; - 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 CHALLENGE_CACHE = new ConcurrentHashMap<>(); private ChallengeParameters challenge; private final boolean disableChallengeResourceVerification; @@ -68,16 +73,17 @@ private static Map extractChallengeAttributes(String authenticat return Collections.emptyMap(); } - authenticateHeader = - authenticateHeader.toLowerCase(Locale.ROOT).replace(authChallengePrefix.toLowerCase(Locale.ROOT), ""); - - String[] attributes = authenticateHeader.split(", "); + String[] attributes = authenticateHeader + .replace("\"", "") + .substring(authChallengePrefix.length()) + .split(","); Map attributeMap = new HashMap<>(); for (String pair : attributes) { - String[] keyValue = pair.split("="); + // Using trim is ugly, but we need it here because currently the 'claims' attribute comes after two spaces. + String[] keyValue = pair.trim().split("=", 2); - attributeMap.put(keyValue[0].replaceAll("\"", ""), keyValue[1].replaceAll("\"", "")); + attributeMap.put(keyValue[0], keyValue[1]); } return attributeMap; @@ -103,31 +109,31 @@ public Mono authorizeRequest(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); 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. + // 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.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((Flux) null); } } @@ -146,7 +152,7 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((Flux) contentOptional.get()); - request.setHeader(CONTENT_LENGTH_HEADER, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); @@ -201,7 +207,23 @@ public Mono authorizeRequestOnChallenge(HttpPipelineCallContext context TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims( + new String(Base64Util.decodeString(claims), StandardCharsets.UTF_8)); + } + } + } return setAuthorizationHeader(context, tokenRequestContext) .then(Mono.just(true)); @@ -214,38 +236,38 @@ public void authorizeRequestSync(HttpPipelineCallContext context) { // If this policy doesn't have challenge parameters cached try to get it from the static challenge cache. if (this.challenge == null) { - String authority = getRequestAuthority(request); - this.challenge = CHALLENGE_CACHE.get(authority); + this.challenge = CHALLENGE_CACHE.get(getRequestAuthority(request)); } if (this.challenge != null) { // We fetched the challenge from the cache, but we have not initialized the scopes in the base yet. TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); setAuthorizationHeaderSync(context, tokenRequestContext); + return; } // 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. + // 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.getBodyAsBinaryData() != null) { context.setData(KEY_VAULT_STASHED_CONTENT_KEY, request.getBodyAsBinaryData()); context.setData(KEY_VAULT_STASHED_CONTENT_LENGTH_KEY, - request.getHeaders().getValue(CONTENT_LENGTH_HEADER)); - request.setHeader(CONTENT_LENGTH_HEADER, "0"); + request.getHeaders().getValue(CONTENT_LENGTH)); + request.setHeader(CONTENT_LENGTH, "0"); request.setBody((BinaryData) null); } } } - @SuppressWarnings("unchecked") @Override public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) { HttpRequest request = context.getHttpRequest(); @@ -254,7 +276,7 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, if (request.getBody() == null && contentOptional.isPresent() && contentLengthOptional.isPresent()) { request.setBody((BinaryData) (contentOptional.get())); - request.setHeader(CONTENT_LENGTH_HEADER, (String) contentLengthOptional.get()); + request.setHeader(CONTENT_LENGTH, (String) contentLengthOptional.get()); } String authority = getRequestAuthority(request); @@ -309,12 +331,131 @@ public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, TokenRequestContext tokenRequestContext = new TokenRequestContext() .addScopes(this.challenge.getScopes()) - .setTenantId(this.challenge.getTenantId()); + .setTenantId(this.challenge.getTenantId()) + .setCaeEnabled(true); + + String error = challengeAttributes.get("error"); + + if (error != null) { + LOGGER.verbose("The challenge response contained an error: {}", error); + + if ("insufficient_claims".equalsIgnoreCase(error)) { + String claims = challengeAttributes.get("claims"); + + if (claims != null) { + tokenRequestContext.setClaims(new String(Base64Util.decodeString(claims))); + } + } + } setAuthorizationHeaderSync(context, tokenRequestContext); + return true; } + @Override + public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + return Mono.error(new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return authorizeRequest(context).then(Mono.defer(next::process)).flatMap(httpResponse -> { + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallenge(context, httpResponse, nextPolicy); + } + + return Mono.just(httpResponse); + }); + } + + @Override + public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) { + if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) { + throw LOGGER.logExceptionAsError( + new RuntimeException("Token credentials require a URL using the HTTPS protocol scheme.")); + } + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + + authorizeRequestSync(context); + + HttpResponse httpResponse = next.processSync(); + String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (httpResponse.getStatusCode() == 401 && authHeader != null) { + return handleChallengeSync(context, httpResponse, nextPolicy); + } + + return httpResponse; + } + + private Mono handleChallenge(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextPolicy next) { + return authorizeRequestOnChallenge(context, httpResponse).flatMap(authorized -> { + if (authorized) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextPolicy nextPolicy = next.clone(); + + return next.process().flatMap(newResponse -> { + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallenge(context, newResponse, nextPolicy); + } else { + return Mono.just(newResponse); + } + }); + } + + return Mono.just(httpResponse); + }); + } + + private HttpResponse handleChallengeSync(HttpPipelineCallContext context, HttpResponse httpResponse, + HttpPipelineNextSyncPolicy next) { + if (authorizeRequestOnChallengeSync(context, httpResponse)) { + // The body needs to be closed or read to the end to release the connection. + httpResponse.close(); + + HttpPipelineNextSyncPolicy nextPolicy = next.clone(); + HttpResponse newResponse = next.processSync(); + String authHeader = newResponse.getHeaderValue(WWW_AUTHENTICATE); + + if (newResponse.getStatusCode() == 401 && authHeader != null && isClaimsPresent(newResponse) + && !isClaimsPresent(httpResponse)) { + + return handleChallengeSync(context, newResponse, nextPolicy); + } + + return newResponse; + } + + return httpResponse; + } + + private boolean isClaimsPresent(HttpResponse httpResponse) { + Map challengeAttributes = + extractChallengeAttributes(httpResponse.getHeaderValue(WWW_AUTHENTICATE), BEARER_TOKEN_PREFIX); + + String error = challengeAttributes.get("error"); + + if (error != null) { + String base64Claims = challengeAttributes.get("claims"); + + return "insufficient_claims".equalsIgnoreCase(error) && base64Claims != null; + } + + return false; + } + private static class ChallengeParameters { private final URI authorizationUri; private final String tenantId; diff --git a/sdk/keyvault/azure-security-keyvault-secrets/src/test/java/com/azure/security/keyvault/secrets/KeyVaultCredentialPolicyTest.java b/sdk/keyvault/azure-security-keyvault-secrets/src/test/java/com/azure/security/keyvault/secrets/KeyVaultCredentialPolicyTest.java index b70ced39fe147..2ec6839db94e5 100644 --- a/sdk/keyvault/azure-security-keyvault-secrets/src/test/java/com/azure/security/keyvault/secrets/KeyVaultCredentialPolicyTest.java +++ b/sdk/keyvault/azure-security-keyvault-secrets/src/test/java/com/azure/security/keyvault/secrets/KeyVaultCredentialPolicyTest.java @@ -3,7 +3,9 @@ package com.azure.security.keyvault.secrets; -import com.azure.core.credential.BasicAuthenticationCredential; +import com.azure.core.credential.AccessToken; +import com.azure.core.credential.TokenCredential; +import com.azure.core.credential.TokenRequestContext; import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; @@ -15,6 +17,7 @@ import com.azure.core.test.SyncAsyncExtension; import com.azure.core.test.annotation.SyncAsyncTest; import com.azure.core.test.http.MockHttpResponse; +import com.azure.core.util.Base64Util; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.security.keyvault.secrets.implementation.KeyVaultCredentialPolicy; @@ -29,11 +32,22 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import java.util.stream.Stream; +import static com.azure.core.http.HttpHeaderName.AUTHORIZATION; +import static com.azure.core.util.CoreUtils.isNullOrEmpty; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -43,18 +57,33 @@ public class KeyVaultCredentialPolicyTest { private static final String AUTHENTICATE_HEADER = "Bearer authorization=\"https://login.windows.net/72f988bf-86f1-41af-91ab-2d7cd022db57\", " + "resource=\"https://vault.azure.net\""; + private static final String AUTHENTICATE_HEADER_WITH_CLAIMS = + "Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", " + + "error=\"insufficient_claims\", " + + "claims=\"eyJhY2Nlc3NfdG9rZW4iOnsiYWNycyI6eyJlc3NlbnRpYWwiOnRydWUsInZhbHVlIjoiY3AxIn19fQ==\""; + private static final String DECODED_CLAIMS = "{\"access_token\":{\"acrs\":{\"essential\":true,\"value\":\"cp1\"}}}"; private static final String BEARER = "Bearer"; private static final String BODY = "this is a sample body"; private static final Flux BODY_FLUX = Flux.defer(() -> Flux.fromStream(Stream.of(BODY.split("")).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8))))); - private BasicAuthenticationCredential credential; + private static final String FAKE_ENCODED_CREDENTIAL = + Base64Util.encodeToString("user:fakePasswordPlaceholder".getBytes(StandardCharsets.UTF_8)); + private static final List> BASE_ASSERTIONS = Arrays.asList( + tokenRequestContext -> !tokenRequestContext.getScopes().isEmpty(), + tokenRequestContext -> !isNullOrEmpty(tokenRequestContext.getTenantId()), + TokenRequestContext::isCaeEnabled); + + private HttpResponse simpleResponse; + private HttpResponse unauthorizedHttpResponseWithWrongStatusCode; private HttpResponse unauthorizedHttpResponseWithHeader; private HttpResponse unauthorizedHttpResponseWithoutHeader; + private HttpResponse unauthorizedHttpResponseWithHeaderAndClaims; private HttpPipelineCallContext callContext; private HttpPipelineCallContext differentScopeContext; private HttpPipelineCallContext testContext; private HttpPipelineCallContext bodyContext; private HttpPipelineCallContext bodyFluxContext; + private TokenCredential credential; private static HttpPipelineCallContext createCallContext(HttpRequest request, Context context) { AtomicReference callContextReference = new AtomicReference<>(); @@ -62,12 +91,13 @@ private static HttpPipelineCallContext createCallContext(HttpRequest request, Co HttpPipeline callContextCreator = new HttpPipelineBuilder() .policies((callContext, next) -> { callContextReference.set(callContext); + return next.process(); }) .httpClient(ignored -> Mono.empty()) .build(); - callContextCreator.send(request, context).block(); + callContextCreator.sendSync(request, context); return callContextReference.get(); } @@ -83,21 +113,37 @@ public void setup() { Context bodyFluxContextContext = new Context("KeyVaultCredentialPolicyStashedBody", BODY_FLUX) .addData("KeyVaultCredentialPolicyStashedContentLength", "21"); + MockHttpResponse simpleResponse = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 200); + + MockHttpResponse unauthorizedResponseWithWrongStatusCode = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + + MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401); + MockHttpResponse unauthorizedResponseWithHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500, + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER)); - MockHttpResponse unauthorizedResponseWithoutHeader = new MockHttpResponse( - new HttpRequest(HttpMethod.GET, "https://azure.com"), 500); + MockHttpResponse unauthorizedResponseWithHeaderAndClaims = new MockHttpResponse( + new HttpRequest(HttpMethod.GET, "https://azure.com"), 401, + new HttpHeaders().set(HttpHeaderName.WWW_AUTHENTICATE, AUTHENTICATE_HEADER_WITH_CLAIMS)); + this.simpleResponse = simpleResponse; + this.unauthorizedHttpResponseWithWrongStatusCode = unauthorizedResponseWithWrongStatusCode; this.unauthorizedHttpResponseWithHeader = unauthorizedResponseWithHeader; this.unauthorizedHttpResponseWithoutHeader = unauthorizedResponseWithoutHeader; + this.unauthorizedHttpResponseWithHeaderAndClaims = unauthorizedResponseWithHeaderAndClaims; this.callContext = createCallContext(request, Context.NONE); this.differentScopeContext = createCallContext(requestWithDifferentScope, Context.NONE); - this.credential = new BasicAuthenticationCredential("user", "fakePasswordPlaceholder"); this.testContext = createCallContext(request, Context.NONE); this.bodyContext = createCallContext(request, bodyContextContext); this.bodyFluxContext = createCallContext(request, bodyFluxContextContext); + // Can't use BasicAuthenticationCredential until the following PR is merged: + // https://github.com/Azure/azure-sdk-for-java/pull/42238 + this.credential = tokenRequestContext -> + Mono.fromCallable(() -> new AccessToken(FAKE_ENCODED_CREDENTIAL, OffsetDateTime.MAX.minusYears(1))); } @AfterEach @@ -105,19 +151,57 @@ public void cleanup() { KeyVaultCredentialPolicy.clearCache(); } + @SyncAsyncTest + public void onNon401ErrorResponse() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithWrongStatusCode)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + + @SyncAsyncTest + public void on401UnauthorizedResponseWithHeader() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(unauthorizedHttpResponseWithHeader)) + .build(); + + SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext()) + ); + + assertNotNull(this.callContext.getHttpRequest().getHeaders().get(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onChallengeCredentialPolicy() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); - String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } @@ -126,16 +210,125 @@ public void onChallengeCredentialPolicy() { public void onAuthorizeRequestChallengeCachePresent() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + StepVerifier.create(onChallengeAndClearCache(policy, this.callContext, unauthorizedHttpResponseWithHeader) // Challenge cache created + .then(policy.authorizeRequest(this.testContext))) // Challenge cache used + .verifyComplete(); + + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + assertFalse(tokenValue.isEmpty()); + assertTrue(tokenValue.startsWith(BEARER)); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Challenge cache created - onChallenge(policy, this.callContext, unauthorizedHttpResponseWithHeader).block(); + onChallengeAndClearCacheSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); // Challenge cache used - policy.authorizeRequest(this.testContext).block(); + policy.authorizeRequestSync(this.testContext); - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); } + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeader) + .flatMap(authorized -> { + if (authorized) { + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + return policy.authorizeRequestOnChallenge(this.callContext, // Challenge with claims received + this.unauthorizedHttpResponseWithHeaderAndClaims) + .map(ignored -> firstToken); + } else { + return Mono.just(""); + } + })) + .assertNext(firstToken -> { + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaims() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + StepVerifier.create(policy.authorizeRequestOnChallenge(this.callContext, // Challenge cache created + this.unauthorizedHttpResponseWithHeaderAndClaims)) + .assertNext(result -> { + assertFalse(result); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + }) + .verifyComplete(); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeCachePresentWithClaimsSync() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + // Challenge cache created + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, this.unauthorizedHttpResponseWithHeader)); + + String firstToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(firstToken.isEmpty()); + assertTrue(firstToken.startsWith(BEARER)); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // Challenge with claims received + assertTrue(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + + String newToken = this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + assertFalse(newToken.isEmpty()); + assertTrue(newToken.startsWith(BEARER)); + + assertNotEquals(firstToken, newToken); + + KeyVaultCredentialPolicy.clearCache(); + } + + @Test + public void onAuthorizeRequestChallengeNoCachePresentWithClaimsSync() { + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + + // Challenge with claims received + assertFalse(policy.authorizeRequestOnChallengeSync(this.callContext, + this.unauthorizedHttpResponseWithHeaderAndClaims)); + assertNull(this.testContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); + + KeyVaultCredentialPolicy.clearCache(); + } + @SyncAsyncTest public void onAuthorizeRequestNoCache() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); @@ -146,7 +339,7 @@ public void onAuthorizeRequestNoCache() { () -> policy.authorizeRequest(this.callContext) ); - assertNull(this.callContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION)); + assertNull(this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION)); } @SyncAsyncTest @@ -154,21 +347,21 @@ public void testSetContentLengthHeader() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.bodyContext, this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.bodyFluxContext, this.unauthorizedHttpResponseWithHeader) ); // Validate that the onChallengeSync ran successfully. assertTrue(onChallenge); HttpHeaders headers = this.bodyFluxContext.getHttpRequest().getHeaders(); - String tokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String tokenValue = headers.getValue(AUTHORIZATION); assertFalse(tokenValue.isEmpty()); assertTrue(tokenValue.startsWith(BEARER)); assertEquals("21", headers.getValue(HttpHeaderName.CONTENT_LENGTH)); HttpHeaders syncHeaders = this.bodyContext.getHttpRequest().getHeaders(); - String syncTokenValue = headers.getValue(HttpHeaderName.AUTHORIZATION); + String syncTokenValue = headers.getValue(AUTHORIZATION); assertFalse(syncTokenValue.isEmpty()); assertTrue(syncTokenValue.startsWith(BEARER)); assertEquals("21", syncHeaders.getValue(HttpHeaderName.CONTENT_LENGTH)); @@ -179,8 +372,8 @@ public void onAuthorizeRequestNoScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), - () -> onChallenge(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) + () -> onChallengeAndClearCacheSync(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader), + () -> onChallengeAndClearCache(policy, this.callContext, this.unauthorizedHttpResponseWithoutHeader) ); assertFalse(onChallenge); @@ -191,9 +384,11 @@ public void onAuthorizeRequestDifferentScope() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); assertThrows(RuntimeException.class, - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)); + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)); - StepVerifier.create(onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader)) + StepVerifier.create(onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader)) .verifyErrorMessage("The challenge resource 'https://vault.azure.net/.default' does not match the " + "requested domain. If you wish to disable this check for your client, pass 'true' to the " + "SecretClientBuilder.disableChallengeResourceVerification() method when building it. See " @@ -205,38 +400,290 @@ public void onAuthorizeRequestDifferentScopeVerifyFalse() { KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, true); boolean onChallenge = SyncAsyncExtension.execute( - () -> onChallengeSync(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader), - () -> onChallenge(policy, this.differentScopeContext, this.unauthorizedHttpResponseWithHeader) + () -> onChallengeAndClearCacheSync(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader), + () -> onChallengeAndClearCache(policy, this.differentScopeContext, + this.unauthorizedHttpResponseWithHeader) ); assertTrue(onChallenge); } - @Test - public void onAuthorizeRequestChallengeCachePresentSync() { - KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(this.credential, false); + // Normal flow: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void processMultipleResponses() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); - // Challenge cache created - onChallengeSync(policy, this.callContext, unauthorizedHttpResponseWithHeader); - // Challenge cache used - policy.authorizeRequestSync(this.testContext); + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. - String tokenValue = this.testContext.getHttpRequest().getHeaders().getValue(HttpHeaderName.AUTHORIZATION); - assertFalse(tokenValue.isEmpty()); - assertTrue(tokenValue.startsWith(BEARER)); + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request. + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + // On receiving an unauthorized response with claims, the token should be updated and a new attempt to make the + // original request should be made. + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the response with claims and set on the request. + assertNotNull(newToken); + // The token was updated. + assertNotEquals(firstToken, newToken); + // A subsequent request was successful. + assertEquals(simpleResponse, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized with claims + @SyncAsyncTest + public void processConsecutiveResponsesWithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response with claims is received, it shall be returned as is. + unauthorizedHttpResponseWithHeaderAndClaims, + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeaderAndClaims, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 200 OK -> 401 Unauthorized with claims -> 401 Unauthorized + @SyncAsyncTest + public void process401WithoutClaimsAfter401WithClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)) + .addAssertion(tokenRequestContext -> tokenRequestContext.getClaims() == null); + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + simpleResponse, + unauthorizedHttpResponseWithHeaderAndClaims, + // If a second consecutive unauthorized response is received, it shall be returned as is. + unauthorizedHttpResponseWithHeader + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String firstToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first response was unauthorized and a token was set on the request + assertNotNull(firstToken); + // On a second attempt, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + testCredential.replaceAssertion(tokenRequestContext -> + DECODED_CLAIMS.equals(tokenRequestContext.getClaims()), 3); + + HttpResponse newResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // A new token was fetched using the first response with claims and set on the request + assertNotEquals(firstToken, newToken); + // A subsequent request was unsuccessful. + assertEquals(unauthorizedHttpResponseWithHeader, newResponse); + + KeyVaultCredentialPolicy.clearCache(); + } + + // Edge case: 401 Unauthorized -> 401 Unauthorized with claims -> 200 OK + @SyncAsyncTest + public void process401WithClaimsAfter401WithoutClaims() { + MutableTestCredential testCredential = new MutableTestCredential(new ArrayList<>(BASE_ASSERTIONS)); + final String[] firstToken = new String[1]; + + testCredential.addAssertion(tokenRequestContext -> { + // This will ensure that that the first request does not contains claims, but the second does after + // receiving a 401 response with a challenge with claims. + testCredential.replaceAssertion(anotherTokenRequestContext -> + DECODED_CLAIMS.equals(anotherTokenRequestContext.getClaims()), 3); + + // We will also store the value of the first credential before it changes on a second call + firstToken[0] = Base64Util.encodeToString(testCredential.getCredential().getBytes(StandardCharsets.UTF_8)); + + assertNotNull(firstToken[0]); + + return tokenRequestContext.getClaims() == null; + }); + + HttpResponse[] responses = new HttpResponse[] { + unauthorizedHttpResponseWithHeader, + unauthorizedHttpResponseWithHeaderAndClaims, + simpleResponse + }; + AtomicInteger currentResponse = new AtomicInteger(); + KeyVaultCredentialPolicy policy = new KeyVaultCredentialPolicy(testCredential, false); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(policy) + .httpClient(ignored -> Mono.just(responses[currentResponse.getAndIncrement()])) + .build(); + + // The first request to a Key Vault endpoint without an access token will always return a 401 Unauthorized + // response with a WWW-Authenticate header containing an authentication challenge. + + HttpResponse firstResponse = SyncAsyncExtension.execute( + () -> pipeline.sendSync(this.callContext.getHttpRequest(), this.callContext.getContext()), + () -> pipeline.send(this.callContext.getHttpRequest(), this.callContext.getContext())); + + String newToken = this.callContext.getHttpRequest().getHeaders().getValue(AUTHORIZATION); + + // The first unauthorized response caused a token to be set on the request, then the token was updated on a + // subsequent unauthorized response with claims. + assertNotEquals(firstToken[0], newToken); + // Finally, a successful response was received. + assertEquals(simpleResponse, firstResponse); + + KeyVaultCredentialPolicy.clearCache(); } - private Mono onChallenge(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private Mono onChallengeAndClearCache(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { Mono onChallenge = policy.authorizeRequestOnChallenge(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallenge; } - private boolean onChallengeSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, - HttpResponse unauthorizedHttpResponse) { + private boolean onChallengeAndClearCacheSync(KeyVaultCredentialPolicy policy, HttpPipelineCallContext callContext, + HttpResponse unauthorizedHttpResponse) { boolean onChallengeSync = policy.authorizeRequestOnChallengeSync(callContext, unauthorizedHttpResponse); + KeyVaultCredentialPolicy.clearCache(); + return onChallengeSync; } + + private static final class MutableTestCredential implements TokenCredential { + private String credential; + private List> assertions; + + private MutableTestCredential(List> assertions) { + this.credential = new Random().toString(); + this.assertions = assertions; + } + + /** + * @throws RuntimeException if any of the assertions fail. + */ + @Override + public Mono getToken(TokenRequestContext requestContext) { + if (requestContext.isCaeEnabled() && requestContext.getClaims() != null) { + credential = new Random().toString(); + } + + String encodedCredential = Base64Util.encodeToString(credential.getBytes(StandardCharsets.UTF_8)); + + for (int i = 0; i < assertions.size(); i++) { + if (!assertions.get(i).apply(requestContext)) { + return Mono.error(new RuntimeException(String.format("Assertion number %d failed", i))); + } + } + + return Mono.fromCallable(() -> new AccessToken(encodedCredential, OffsetDateTime.MAX.minusYears(1))); + } + + private MutableTestCredential setAssertions(List> assertions) { + this.assertions = assertions; + + return this; + } + + private MutableTestCredential addAssertion(Function assertion) { + assertions.add(assertion); + + return this; + } + + private MutableTestCredential replaceAssertion(Function assertion, int index) { + assertions.set(index, assertion); + + return this; + } + + private String getCredential() { + return this.credential; + } + } }