Default CAE challenge handling in BearerTokenAuthenticationPolicy (#4…

…2471)

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

@@ -13,9 +13,13 @@
 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.nio.charset.StandardCharsets;
+import java.util.Base64;
 import java.util.Objects;
 
 /**
@@ -75,7 +79,8 @@ 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 +89,53 @@ 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 +224,25 @@ 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");
+
+        if (!CoreUtils.isNullOrEmpty(encodedClaims)) {
+            try {
+                decodedClaims = new String(Base64.getDecoder().decode(encodedClaims), StandardCharsets.UTF_8);
+            } 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,121 @@
+// 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;
+
+/**
+ * Parses Authorization challenges from the {@link HttpResponse}.
+ */
+public final class AuthorizationChallengeParser {
+
+    /**
+     * Creates an instance of the AuthorizationChallengeParser.
+     */
+    private AuthorizationChallengeParser() {
+    }
+
+    /**
+     * 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 challenge = response.getHeaderValue(HttpHeaderName.WWW_AUTHENTICATE);
+
+        String parameters = getChallengeParametersForScheme(challenge, "Bearer");
+        String error = getChallengeParameterValue(parameters, "error");
+        String claims = getChallengeParameterValue(parameters, "claims");
+        return !CoreUtils.isNullOrEmpty(claims) && "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);
+        String parameters = getChallengeParametersForScheme(challenge, challengeScheme);
+        return getChallengeParameterValue(parameters, parameter);
+    }
+
+    /**
+     * Gets the set of challenge parameters for the specified challenge scheme.
+     * @param challenge The challenge to parse.
+     * @param challengeScheme The challenge scheme to extract parameters for.
+     * @return The extracted challenge parameters for the specified challenge scheme.
+     */
+    private static String getChallengeParametersForScheme(String challenge, String challengeScheme) {
+        if (CoreUtils.isNullOrEmpty(challenge)) {
+            return null;
+        }
+
+        int schemeIndex = -1;
+        int length = challenge.length();
+        int schemeLength = challengeScheme.length();
+
+        for (int i = 0; i <= length - schemeLength - 1; i++) {
+            // Check if the scheme matches and is followed by a space
+            if (challenge.startsWith(challengeScheme, i)
+                && (i + schemeLength < length)
+                && challenge.charAt(i + schemeLength) == ' ') {
+                schemeIndex = i;
+                break;
+            }
+        }
+
+        if (schemeIndex == -1) {
+            return null; // Scheme not found
+        }
+
+        int startIndex = schemeIndex + challengeScheme.length();
+        int endIndex = challenge.length();
+
+        // Skip whitespace after the scheme to avoid unnecessary trim
+        while (startIndex < endIndex && Character.isWhitespace(challenge.charAt(startIndex))) {
+            startIndex++;
+        }
+
+        // Skip trailing whitespace
+        while (endIndex > startIndex && Character.isWhitespace(challenge.charAt(endIndex - 1))) {
+            endIndex--;
+        }
+
+        return startIndex < endIndex ? challenge.substring(startIndex, endIndex) : null;
+    }
+
+    /**
+     * Gets the specified challenge parameter from the challenge.
+     * @param parameters The challenge parameters to parse.
+     * @param parameter The parameter to extract.
+     * @return The extracted value of the challenge parameter.
+     */
+    private static String getChallengeParameterValue(String parameters, String parameter) {
+        if (CoreUtils.isNullOrEmpty(parameters)) {
+            return null;
+        }
+
+        String[] paramPairs = parameters.split(",", -1);
+        for (String pair : paramPairs) {
+            int equalsIndex = pair.indexOf('=');
+            if (equalsIndex != -1) {
+                String key = pair.substring(0, equalsIndex).trim();
+
+                if (key.equals(parameter)) {
+                    String value = pair.substring(equalsIndex + 1).replace("\"", "").trim();
+                    return value;
+                }
+            }
+        }
+        return null;
+    }
+}

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

@@ -0,0 +1,132 @@
+// 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.implementation.http.policy.AuthorizationChallengeParser;
+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());
+
+        if (expectedClaims != null) {
+            String actualEncodedClaims = AuthorizationChallengeParser.getChallengeParameterFromResponse(
+                new MockHttpResponse(null, 401, new HttpHeaders().add(HttpHeaderName.WWW_AUTHENTICATE, challenge)),
+                "Bearer", "claims");
+            assertEquals(encodedClaims, actualEncodedClaims);
+        }
+    }
+
+    @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());
+
+        if (expectedClaims != null) {
+            String actualEncodedClaims = AuthorizationChallengeParser.getChallengeParameterFromResponse(
+                new MockHttpResponse(null, 401, new HttpHeaders().add(HttpHeaderName.WWW_AUTHENTICATE, challenge)),
+                "Bearer", "claims");
+            assertEquals(encodedClaims, actualEncodedClaims);
+        }
+    }
+
+    // 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
+            Arguments.of("Bearer claims=\"\" error=\"insufficient_claims\"", 401, null, ""), // empty claims
+            Arguments.of("Bearer error=\"insufficient_claims\"", 401, null, "") // missing claims
+        );
+    }
+}