Default CAE challenge handling in BearerTokenAuthenticationPolicy

This adds handling CAE challenges by default in our policy.

Fixes Azure#41984

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/BearerTokenAuthenticationPolicy.java

@@ -13,9 +13,12 @@
 import com.azure.core.http.HttpPipelineNextSyncPolicy;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.implementation.AccessTokenCache;
+import com.azure.core.implementation.http.policy.AuthorizationChallengeParser;
+import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 import reactor.core.publisher.Mono;
 
+import java.util.Base64;
 import java.util.Objects;
 
 /**
@@ -75,7 +78,7 @@ public Mono<Void> authorizeRequest(HttpPipelineCallContext context) {
         if (this.scopes == null) {
             return Mono.empty();
         }
-        return setAuthorizationHeaderHelper(context, new TokenRequestContext().addScopes(this.scopes), false);
+        return setAuthorizationHeaderHelper(context, new TokenRequestContext().addScopes(this.scopes).setCaeEnabled(true), false);
     }
 
     /**
@@ -84,32 +87,52 @@ public Mono<Void> authorizeRequest(HttpPipelineCallContext context) {
      * @param context The request context.
      */
     public void authorizeRequestSync(HttpPipelineCallContext context) {
-        setAuthorizationHeaderHelperSync(context, new TokenRequestContext().addScopes(scopes), false);
+        setAuthorizationHeaderHelperSync(context, new TokenRequestContext().addScopes(scopes).setCaeEnabled(true), false);
     }
 
     /**
      * Handles the authentication challenge in the event a 401 response with a WWW-Authenticate authentication challenge
      * header is received after the initial request and returns appropriate {@link TokenRequestContext} to be used for
      * re-authentication.
+     * <p>
+     * The default implementation will attempt to handle Continuous Access Evaluation (CAE) challenges.
+     * </p>
      *
      * @param context The request context.
      * @param response The Http Response containing the authentication challenge header.
      * @return A {@link Mono} containing {@link TokenRequestContext}
      */
     public Mono<Boolean> authorizeRequestOnChallenge(HttpPipelineCallContext context, HttpResponse response) {
+        if (AuthorizationChallengeParser.isCaeClaimsChallenge(response)) {
+            TokenRequestContext tokenRequestContext = getTokenRequestContextForCaeChallenge(response);
+            if (tokenRequestContext != null) {
+                return setAuthorizationHeader(context, tokenRequestContext).then(Mono.just(true));
+            }
+        }
         return Mono.just(false);
     }
 
     /**
      * Handles the authentication challenge in the event a 401 response with a WWW-Authenticate authentication challenge
      * header is received after the initial request and returns appropriate {@link TokenRequestContext} to be used for
      * re-authentication.
+     * <p>
+     * The default implementation will attempt to handle Continuous Access Evaluation (CAE) challenges.
+     * </p>
      *
      * @param context The request context.
      * @param response The Http Response containing the authentication challenge header.
      * @return A boolean indicating if containing the {@link TokenRequestContext} for re-authentication
      */
     public boolean authorizeRequestOnChallengeSync(HttpPipelineCallContext context, HttpResponse response) {
+        if (AuthorizationChallengeParser.isCaeClaimsChallenge(response)) {
+            TokenRequestContext tokenRequestContext = getTokenRequestContextForCaeChallenge(response);
+            if (tokenRequestContext != null) {
+                setAuthorizationHeaderSync(context, tokenRequestContext);
+                return true;
+            }
+        }
+
         return false;
     }
 
@@ -198,4 +221,24 @@ private void setAuthorizationHeaderHelperSync(HttpPipelineCallContext context,
     private static void setAuthorizationHeader(HttpHeaders headers, String token) {
         headers.set(HttpHeaderName.AUTHORIZATION, BEARER + " " + token);
     }
+
+    private TokenRequestContext getTokenRequestContextForCaeChallenge(HttpResponse response) {
+        String decodedClaims = null;
+        String encodedClaims = AuthorizationChallengeParser.getChallengeParameterFromResponse(response, "Bearer", "claims");
+
+        try {
+            if (!CoreUtils.isNullOrEmpty(encodedClaims)) {
+                decodedClaims = new String(Base64.getDecoder().decode(encodedClaims));
+            }
+        } catch (IllegalArgumentException e) {
+            // We don't want to throw here, but we want to log this for future incident investigation.
+            LOGGER.warning("Failed to decode the claims from the CAE challenge. Encoded claims" + encodedClaims);
+        }
+
+        if (decodedClaims == null) {
+            return null;
+        }
+
+        return new TokenRequestContext().setClaims(decodedClaims).addScopes(scopes).setCaeEnabled(true);
+    }
 }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/policy/AuthorizationChallengeParser.java

@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.implementation.http.policy;
+
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.util.CoreUtils;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * Parses Authorization challenges from the {@link HttpResponse}.
+ */
+public class AuthorizationChallengeParser {
+
+    private static final Pattern CHALLENGE_PATTERN = Pattern.compile("(\\w+) ((?:\\w+=\"[^\"]*\",?\\s*)+)");
+    private static final Pattern CHALLENGE_PARAMS_PATTERN = Pattern.compile("(\\w+)=\"([^\"]*)\"");
+
+    /**
+     * Examines a {@link HttpResponse} to see if it is a CAE challenge.
+     * @param response The {@link HttpResponse} to examine.
+     * @return True if the response is a CAE challenge, false otherwise.
+     */
+    public static boolean isCaeClaimsChallenge(HttpResponse response) {
+        String error = getChallengeParameterFromResponse(response, "Bearer", "error");
+        return "insufficient_claims".equals(error);
+    }
+
+    /**
+     * Gets the specified challenge parameter from the challenge response.
+     *
+     * @param response the Http response with auth challenge
+     * @param challengeScheme the challenge scheme to be checked
+     * @param parameter the challenge parameter value to get
+     *
+     * @return the extracted value of the challenge parameter
+     */
+    public static String getChallengeParameterFromResponse(HttpResponse response, String challengeScheme,
+                                                           String parameter) {
+        String challenge = response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE);
+        return getChallengeParameter(challenge, challengeScheme, parameter);
+    }
+
+    /**
+     * Gets the specified challenge parameter from the challenge.
+     * @param challenge The challenge string to parse.
+     * @param challengeScheme The requested scheme (e.g. "Bearer" or "PoP")
+     * @param parameter The parameter to extract.
+     * @return The extracted value of the challenge parameter.
+     */
+    private static String getChallengeParameter(String challenge, String challengeScheme, String parameter) {
+        if (CoreUtils.isNullOrEmpty(challenge)) {
+            return null;
+        }
+
+        Matcher challengeMatch = CHALLENGE_PATTERN.matcher(challenge);
+        while (challengeMatch.find()) {
+            if (challengeMatch.group(1).equals(challengeScheme)) {
+                Matcher paramsMatch = CHALLENGE_PARAMS_PATTERN.matcher(challengeMatch.group(2));
+                while (paramsMatch.find()) {
+                    if (parameter.equals(paramsMatch.group(1))) {
+                        return paramsMatch.group(2);
+                    }
+                }
+            }
+        }
+        return null;
+    }
+}

sdk/core/azure-core/src/test/java/com/azure/core/http/policy/BearerTokenAuthenticationPolicyTests.java

@@ -0,0 +1,111 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http.policy;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.MockHttpResponse;
+import com.azure.core.util.Context;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.time.OffsetDateTime;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+
+public class BearerTokenAuthenticationPolicyTests {
+
+    @ParameterizedTest
+    @MethodSource("caeTestArguments")
+    public void testDefaultCae(String challenge, int expectedStatusCode, String expectedClaims, String encodedClaims) {
+        AtomicReference<String> claims = new AtomicReference<>();
+        AtomicInteger callCount = new AtomicInteger();
+        TokenCredential credential = getCaeTokenCredential(claims, callCount);
+        BearerTokenAuthenticationPolicy policy = new BearerTokenAuthenticationPolicy(credential, "scope");
+        HttpClient client = getCaeHttpClient(challenge, callCount);
+
+        HttpPipeline pipeline = new HttpPipelineBuilder().policies(policy).httpClient(client).build();
+        StepVerifier.create(pipeline.send(new HttpRequest(HttpMethod.GET, "https://localhost")))
+            .assertNext(response -> assertEquals(expectedStatusCode, response.getStatusCode()))
+            .verifyComplete();
+        assertEquals(expectedClaims, claims.get());
+    }
+
+    @ParameterizedTest
+    @MethodSource("caeTestArguments")
+    public void testDefaultCaeSync(String challenge, int expectedStatusCode, String expectedClaims, String encodedClaims) {
+        AtomicReference<String> claims = new AtomicReference<>();
+        AtomicInteger callCount = new AtomicInteger();
+
+        TokenCredential credential = getCaeTokenCredential(claims, callCount);
+        BearerTokenAuthenticationPolicy policy = new BearerTokenAuthenticationPolicy(credential, "scope");
+        HttpClient client = getCaeHttpClient(challenge, callCount);
+        HttpPipeline pipeline = new HttpPipelineBuilder().policies(policy).httpClient(client).build();
+
+        try (HttpResponse response = pipeline.sendSync(new HttpRequest(HttpMethod.GET, "https://localhost"), Context.NONE)) {
+            assertEquals(expectedStatusCode, response.getStatusCode());
+        }
+        assertEquals(expectedClaims, claims.get());
+    }
+
+    // A fake token credential that lets us keep track of what got parsed out of a CAE claim for assertion.
+    private static TokenCredential getCaeTokenCredential(AtomicReference<String> claims, AtomicInteger callCount) {
+        return request -> {
+            claims.set(request.getClaims());
+            assertTrue(request.isCaeEnabled());
+            callCount.incrementAndGet();
+            return Mono.just(new AccessToken("token", OffsetDateTime.now().plusHours(2)));
+        };
+    }
+
+    // This http client is effectively a state sentinel for how we progressed through the challenge.
+    // If we had a challenge, and it is invalid, the policy stops and returns 401 all the way out.
+    // If the CAE challenge parses properly we will end complete the policy normally and get 200.
+    private static HttpClient getCaeHttpClient(String challenge, AtomicInteger callCount) {
+        return request -> {
+            if (callCount.get() <= 1) {
+                if (challenge == null) {
+                    return Mono.just(new MockHttpResponse(request, 200));
+                }
+                return Mono.just(new MockHttpResponse(request, 401, new HttpHeaders().add(HttpHeaderName.WWW_AUTHENTICATE, challenge)));
+            }
+            return Mono.just(new MockHttpResponse(request, 200));
+        };
+    }
+
+    private static Stream<Arguments> caeTestArguments() {
+        return Stream.of(Arguments.of(null, 200, null, null), // no challenge
+            Arguments.of("Bearer authorization_uri=\"https://login.windows.net/\", error=\"invalid_token\", claims=\"ey==\"", 401, null, "ey=="), // unexpected error value
+            Arguments.of("Bearer claims=\"not base64\", error=\"insufficient_claims\"", 401, null, "not base64"), // parsing error
+            Arguments.of("Bearer realm=\"\", authorization_uri=\"http://localhost\", client_id=\"00000003-0000-0000-c000-000000000000\", error=\"insufficient_claims\", claims=\"ey==\"",
+                    200,
+                    "{",
+                    "ey=="), // more parameters in a different order
+            Arguments.of("Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", error=\"insufficient_claims\", claims=\"eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwidmFsdWUiOiIxNzI2MDc3NTk1In0sInhtc19jYWVlcnJvciI6eyJ2YWx1ZSI6IjEwMDEyIn19fQ==\"",
+                    200,
+                    "{\"access_token\":{\"nbf\":{\"essential\":true,\"value\":\"1726077595\"},\"xms_caeerror\":{\"value\":\"10012\"}}}",
+                    "eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwidmFsdWUiOiIxNzI2MDc3NTk1In0sInhtc19jYWVlcnJvciI6eyJ2YWx1ZSI6IjEwMDEyIn19fQ=="), // standard request
+            Arguments.of("PoP realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", client_id=\"00000003-0000-0000-c000-000000000000\", nonce=\"ey==\", Bearer realm=\"\", authorization_uri=\"https://login.microsoftonline.com/common/oauth2/authorize\", client_id=\"00000003-0000-0000-c000-000000000000\", error_description=\"Continuous access evaluation resulted in challenge with result: InteractionRequired and code: TokenIssuedBeforeRevocationTimestamp\", error=\"insufficient_claims\", claims=\"eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTcyNjI1ODEyMiJ9fX0=\"",
+                    200,
+                    "{\"access_token\":{\"nbf\":{\"essential\":true, \"value\":\"1726258122\"}}}",
+                    "eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTcyNjI1ODEyMiJ9fX0=") // multiple challenges
+        );
+    }
+}