Add request and challenge callbacks to BearerTokenAuthorizationPolicy

[chlowell]

These enable SDK clients to implement custom request authorization and challenge handling (closes #1756).

New API:

• OnRequest trait: defines a callback invoked on each request, responsible for authorizing it before sending
  • Key Vault clients will use this to skip authorization so they can discover auth parameters from a challenge
• OnChallenge trait: defines a callback invoked upon receiving an auth challenge. Implementations are responsible for parsing challenges and authorizing requests and telling the policy whether to retry them
• BearerTokenAuthorizationPolicy builder methods: with_on_request() and with_on_challenge() set callbacks for a policy instance
  • when these aren't set, the policy behaves as it does today: it authorizes each request with its credential and scopes, and returns 401 responses to the pipeline instead of trying to handle their challenges
• Authorizer trait (sealed): allows callbacks to authorize requests and cache tokens without using the policy's credential directly
  • I moved the authNZ logic from send() into a private implementation of this
• Request::body_mut(): lets the policy reset a body stream before retrying after a challenge

[github-actions]

API Change Check

APIView identified API level changes in this PR and created the following API reviews

azure_core
typespec_client_core

[Copilot]

Pull request overview

This PR adds extensible request authorization and authentication challenge handling to BearerTokenAuthorizationPolicy, enabling SDK clients to implement custom authorization and challenge handling logic.

• Introduces three new public traits: OnRequest, OnChallenge, and Authorizer (sealed) to define callbacks for request authorization and challenge handling
• Adds builder methods with_on_request() and with_on_challenge() to configure policy behavior
• Adds Request::body_mut() method to support resetting body streams when retrying after challenges
• Adds WWW_AUTHENTICATE header constant for accessing authentication challenge headers

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
sdk/core/typespec_client_core/src/http/request/mod.rs	Adds body_mut() method to enable mutable access to request body for stream reset operations
sdk/core/typespec_client_core/src/http/headers/common.rs	Adds WWW_AUTHENTICATE header constant following existing header constant patterns
sdk/core/typespec_client_core/CHANGELOG.md	Documents the Request::body_mut() API addition
sdk/core/azure_core/src/http/policies/mod.rs	Exports new public traits (Authorizer, OnChallenge, OnRequest) from bearer_token_policy module
sdk/core/azure_core/src/http/policies/bearer_token_policy.rs	Refactors authorization logic into BearerTokenAuthorizer implementing Authorizer trait; adds OnRequest and OnChallenge trait definitions; implements challenge handling with body stream reset; adds comprehensive test coverage for new functionality
sdk/core/azure_core/CHANGELOG.md	Documents new extensible authorization and challenge handling features

[Copilot]

According to the coding guidelines, changes in sdk/core/typespec_client_core/CHANGELOG.md should be reflected in sdk/core/azure_core/CHANGELOG.md. The addition of Request::body_mut() documented in typespec_client_core/CHANGELOG.md should also be mentioned here since it's a public API change that affects azure_core users.

[heaths]

Copilot is right about the body_mut() call but, Copilot, you're wrong about the direction of changes reflecting in TSS to AC.

[heaths]

Nice, clean design overall but I have a few questions and concerns.

[heaths]

Copilot is right about the body_mut() call but, Copilot, you're wrong about the direction of changes reflecting in TSS to AC.

[heaths]

Doc comments in Rust are like (good) git commits: first line is a summary, blank line, more info. So this should probably be:

Suggested change

	/// Sets a callback for `send` to invoke once on each request it receives, before sending the request. See [`OnRequest`]
	/// for more details. When not set, the policy authorizes each request using the credential and scopes specified to `new`.
	/// Sets a callback for `send` to invoke once on each request it receives, before sending the request.
	///
	/// See [`OnRequest`] for more details. When not set, the policy authorizes each request using the credential and scopes specified to `new`.

[heaths]

The doc comments for with_on_challenge are good by comparison.

[heaths]

Would this let Key Vault, for example, zero out the request body until it has authenticated? That's what we did in .NET by checking if we had an Authorization header or something like that.

[chlowell]

Yes; request is mutable. My prototype for Key Vault has a single type implementing both traits so an instance can hold the challenge request's body. OnRequest stashes the body and OnChallenge restores it.

[heaths]

And a bool should always be sufficient? You can't imagine expanding on that such that maybe we should return a result model that is currently just a bool? If you're reasonably certain, I see no problem with this.

[chlowell]

Yes; attempting to handle a challenge won't produce useful data for the caller. If the callback handled the challenge, the request is ready to go; the policy needs only to send it. If the callback didn't handle the challenge, the policy can't either. And the policy has the same input data--the challenge--as the callback anyway.

But I suppose this could return simply Result<()>. Then Ok means retry the request, and the policy wraps Err with the response and returns that to the pipeline. 🤔 that's simpler and compels the implementer of the callback to write an error message that may be more helpful than the one check_success would otherwise produce. What do you think?

[heaths]

This borderlines on the summary + blank line + description format. Something to consider. Maybe it's fine as-is.

[heaths]

Don't define your own. Just use crate::private::Sealed. Ancestors' private members are visible to all descendants.

[heaths]

What's the purpose of this being a trait? You end up having to pass it around as dyn (or maybe not; impl is better if possible) which is slower to call and this seems like it's in a hot path.

This isn't .NET. We try to avoid traits (similar to .NET's interfaces) as much as possible, unlike .NET. Why not just expose the BearerTokenAuthorizer, especially if it's the only thing that would implement this trait. Seems like an unnecessary abstraction.

[chlowell]

I believe this could be something like

pub type Authorize = dyn for<'a> Fn(
        &'a mut Request,
        &'a [&'a str],
        TokenRequestOptions<'a>,
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
    + Send
    + Sync

instead, if that's preferable. The rationale for the trait aside from making the signature more legible and making it easy for an implementation to hold state is that I expect we'll eventually have at least one more implementation, supporting a different auth scheme for a separate auth policy (TokenRequestOptions will capture the differences between schemes).

[LarryOsterman]

Lines 79 to 96 can be simplified using .then - it alleviates the need for the has_challenge variable at all.

Something like:

let response = if let Some(ref on_challenge) = self.on_challenge {
  response.headers
    .get_optional_str(&WWW_AUTHENTICATE)
    .then(|| {
      let should_retry = on_challenge.on_challenge(ctx, request.self.authorizer.as_ref(), response.headers()).await?;
      if should_retry
         :
         :
    }).or_else(||Ok(response))?;

Pseudo code written in the UX, so take it with a grain of salt.

[chlowell]

It isn't tidier in practice because the signature being and_then<U, F>(self, f: F) -> Option<U> means the closure can't use ?. No problem inlining has_challenge though, and this check can be more terse and a little more efficient by calling get_str instead of get_optional_str

[LarryOsterman]

The idea is that the .and_then closure returns a Result<T> (that's why the Ok(response)) - you then can use ? when it's all done.

Ultimately it's not a huge deal - the code appears to be correct. I'm just a fan of chaining primitives when possible - I find the code slightly easier to understand.

[LarryOsterman]

While the trait can only be implemented within this module, the authorize function should still have a decent doccomment.

[chlowell]

Do you mean it's too brief? What do you suggest adding?

[LarryOsterman]

Arguments mostly - I was thinking of what happens if some service needs its own request/challenge callbacks - they're going to need to access documentation on the Authorizer. And since it's a public trait, it could be used by a service client.