Add native AAD credential callback by ananth7592 · Pull Request #5082 · Azure/azure-sdk-for-rust

ananth7592 · 2026-08-18T16:10:30Z

Add native AAD credential callback

Summary

Adds an asynchronous host-token callback to the Cosmos native wrapper so
language SDKs can supply their own AAD credential without passing a
language-specific object across the C ABI.

The wrapper adapts that callback into azure_core::TokenCredential, caches the
returned token, and performs a single-flight refresh when fewer than two
minutes remain. Master-key authentication is unchanged.

Design

Initial token acquisition

Go owns the language-specific azcore.TokenCredential and registers C
callbacks when it creates the native account reference. On the first Cosmos
operation, Rust has no cached token, so it requests one from Go and waits for
the asynchronous completion.

sequenceDiagram
  participant App as Go Cosmos SDK v2
  participant Credential as Go azcore.TokenCredential
  participant ABI as C ABI callback
  participant Native as Rust CallbackTokenCredential
  participant Driver as Rust Cosmos driver
  participant Cosmos as Cosmos DB

  App->>ABI: Create account with get_token callback and cgo.Handle
  ABI->>Native: Construct azure_core::TokenCredential adapter
  App->>Driver: Submit first Cosmos operation
  Driver->>Native: get_token(Cosmos scope)
  Native->>Native: Cache miss
  Native->>ABI: Invoke get_token(request, completion)
  ABI->>Credential: GetToken asynchronously
  Credential-->>ABI: Access token and expires_on
  ABI-->>Native: Complete exactly once
  Native->>Native: Copy and cache token
  Native-->>Driver: Return AccessToken
  Driver->>Cosmos: Send authorized request
  Cosmos-->>App: Return response

Proactive refresh near expiry

Rust checks the cached token whenever the driver requests authentication. If
two minutes or less remain, Rust starts one refresh and invokes the Go callback
again. Concurrent callers wait for that same refresh rather than starting
additional credential requests.

sequenceDiagram
  participant App as Go Cosmos SDK v2
  participant Credential as Go azcore.TokenCredential
  participant ABI as C ABI callback
  participant Native as Rust CallbackTokenCredential
  participant Driver as Rust Cosmos driver
  participant Cosmos as Cosmos DB

  App->>Driver: Submit later Cosmos operation
  Driver->>Native: get_token(Cosmos scope)
  Native->>Native: Cached token has 2 minutes or less remaining
  Native->>Native: Acquire single-flight refresh lock
  Native->>ABI: Invoke get_token(request, completion)
  ABI->>Credential: GetToken asynchronously
  Credential-->>ABI: Refreshed token and expires_on
  ABI-->>Native: Complete exactly once
  Native->>Native: Replace cached token
  Native-->>Driver: Return refreshed AccessToken
  Driver->>Cosmos: Send authorized request
  Cosmos-->>App: Return response

The new C ABI includes:

cosmos_token_request_t, carrying the requested scope and a completion
callback.
cosmos_token_provider_t, carrying get_token and user_data_free
callbacks.
cosmos_account_ref_with_credential, which creates an account reference
backed by the callback credential.

When get_token returns zero, the host owns completion responsibility and must
invoke the completion callback exactly once. A nonzero return rejects the
request synchronously, and the host must not invoke completion. Token and error
buffers are borrowed only for the completion call; Rust copies them before
returning.

The account reference keeps the provider state alive. user_data_free runs
after the final Rust credential reference is dropped.

Go v2 integration

Go v2 can adapt an azcore.TokenCredential as follows:

Store the Go credential state in runtime/cgo.Handle and pass the handle
value as intptr_t; do not retain a Go pointer in native memory.
Implement get_token and user_data_free C-exported callbacks.
Copy the requested scope before returning from get_token.
Call GetToken asynchronously, then invoke the Rust completion callback
exactly once with either the token and expiry or an error.
Let the Rust adapter cache valid tokens and single-flight refreshes.

The requested Cosmos scope is https://cosmos.azure.com/.default. Ordinary
Cosmos 401 responses are not assumed to mean token expiry and are not
automatically replayed by this change.

The end-to-end Go/cgo proof of concept used to validate this contract is not
included in this Rust-native PR.

Validation

cargo fmt -p azure_data_cosmos_driver_native -- --check
cargo test -p azure_data_cosmos_driver_native --all-features
cargo clippy -p azure_data_cosmos_driver_native --all-features --all-targets -- -D warnings
cSpell on all changed files
Live Go/cgo validation using AzureCLICredential:
- repeated Cosmos item reads reused the cached Rust token;
- a shortened first-token expiry triggered a second Go callback inside the
  two-minute refresh window.

azure-pipelines · 2026-08-18T16:10:42Z

Azure Pipelines: 3 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run.

Adapt host-managed asynchronous token providers into Azure Core's TokenCredential interface so Go v2 and other FFI consumers can use Microsoft Entra authentication without exposing language-specific credential objects across the C ABI. Cache tokens in the native adapter, refresh them proactively, coalesce concurrent refreshes, and define explicit completion and host-state ownership contracts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: df3d3330-e9f9-48b7-bc72-836cb9733f94

Warm-cache reads now go through a shared read lock, matching the pattern in `azure_identity::TokenCache` (sdk/identity/azure_identity/src/cache.rs). Concurrent callers with a fresh cached token no longer serialize behind a Mutex. Single-flight refresh is preserved: the slow path takes an exclusive write lock and double-checks the cache before invoking the host callback, so racing losers become cache hits. The `concurrent_requests_single_flight_refresh` test still asserts a single host callback for 4 concurrent misses. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 92b85280-3c74-4990-ba2b-f69fada72d96

…e_identity Bump TOKEN_REFRESH_SKEW from 2 to 5 minutes to match azure_identity::TokenCache::should_refresh (300s), and drop the post-fetch expired-token rejection. azure_identity performs no post-fetch validation on the token returned by its callback, so this credential enforcing a stricter contract made Cosmos the one credential in the Rust SDK ecosystem that could hard-fail on a short-lived host token. Flip rejects_expired_host_token into returns_short_lived_host_token_verbatim so the deliberate alignment is documented in the test suite. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 92b85280-3c74-4990-ba2b-f69fada72d96

azure-pipelines · 2026-08-19T00:42:25Z

Azure Pipelines: Successfully started running 1 pipeline(s). 2 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run.

Copilot

Pull request overview

Adds native asynchronous AAD credential callbacks for Cosmos language SDK integrations.

Changes:

Adds callback-based token acquisition, caching, and refresh coordination.
Exposes credential construction through the C ABI.
Documents the callback ownership contract and updates dependencies.

Reviewed changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`Cargo.lock`	Records the async-trait dependency.
`sdk/cosmos/azure_data_cosmos_driver/docs/NATIVE_WRAPPER_SPEC.md`	Documents the credential ABI contract.
`sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml`	Adds async-trait.
`sdk/cosmos/azure_data_cosmos_driver_native/build.rs`	Configures credential type names for cbindgen.
`sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h`	Exposes token callbacks and account construction.
`sdk/cosmos/azure_data_cosmos_driver_native/src/account_ref.rs`	Adds callback-credential account references.
`sdk/cosmos/azure_data_cosmos_driver_native/src/credential.rs`	Implements token callbacks, caching, and tests.
`sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs`	Exports the credential module.

Suppressed comments (1)

sdk/cosmos/azure_data_cosmos_driver_native/src/credential.rs:207

The write lock single-flights only while this awaiting future remains alive. Cancellation explicitly drops the driver work (submit.rs:189-190), releasing the guard while the accepted host request must continue (NATIVE_WRAPPER_SPEC.md:1675); the next caller can then start a second host request before the first completes. Store the in-flight refresh independently of any one waiter so cancellation removes only that waiter, not the single-flight state.

        let token = self.request_token(scope).await?;

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

+// sdk/identity/azure_identity/src/cache.rs (300 seconds). Keeping this in sync
+// avoids surprising divergence in when a token is treated as "close enough to
+// expiry" to refresh across the two credential stacks a host can use.
+const TOKEN_REFRESH_SKEW: Duration = Duration::minutes(5);


+        let completion_context =
+            Box::into_raw(Box::new(PendingTokenRequest { sender })).cast::<c_void>();


Kevin Pilch (Pilchie) · 2026-08-19T22:00:38Z

+  /**
+   * Rust completion callback the host invokes exactly once.
+   */
+  cosmos_token_completion_t completion;


Why does this need to be a callback? Couldn't it be just a normal function, and rust maintains a map of context -> outstanding operations?

Kevin Pilch (Pilchie) · 2026-08-19T22:03:13Z

+4. `completion` borrows token and error buffers only for its duration. Rust copies them before returning, destroys `completion_context`, and wakes the pending driver request.
+5. After the final native account/driver credential reference is dropped, Rust invokes `user_data_free` exactly once when it is non-NULL.
+
+A successful completion supplies a non-empty UTF-8 token and a Unix-seconds expiry later than the current time. A failed completion supplies a nonzero status and may supply a UTF-8 error message. The adapter requests exactly one scope (currently `https://cosmos.azure.com/.default`), reuses a cached token while it has more than two minutes remaining, and refreshes proactively inside that window. Cosmos 401 responses do not trigger credential refresh or request replay.


Two minutes seems very short. I thought for .NET we had a custom policy that tries to refresh at 50% of the token lifetime, so that there is time for retries if the token server is down. Would should probably do the same here.

Tomas Varon (tvaron3) · 2026-08-26T20:50:19Z

+        // misses. Double-check the cache in case another task refreshed while
+        // we were waiting on the write lock.
+        let mut cached = self.cached_token.write().await;
+        let now = OffsetDateTime::now_utc();


nit: could we separate out this logic into helper method since its same above

Tomas Varon (tvaron3) · 2026-08-26T21:26:32Z

+        // Slow path: exclusive lock single-flights refresh across concurrent
+        // misses. Double-check the cache in case another task refreshed while
+        // we were waiting on the write lock.
+        let mut cached = self.cached_token.write().await;


These locks would be released when the get token call is cancelled or succeeds. If it is cancellled, the callback in the native sdk layer still continues and the following request would still trigger another refresh. I recommend a test that cancels the initiating caller, then requests another token; asserting only one host callback occurs.

github-actions Bot added the Cosmos The azure_cosmos crate label Aug 18, 2026

ananth7592 changed the base branch from ananth7592-m2-native-driver-pipeline to main August 18, 2026 16:48

ananth7592 force-pushed the ananth7592-cosmos-native-aad-callback branch from 6bc5ff4 to 480063e Compare August 18, 2026 17:44

ananth7592 added 2 commits August 18, 2026 16:07

ananth7592 mentioned this pull request Aug 19, 2026

AsyncHostTokenCredential: no timeout on host callback await #5084

Open

Merge branch 'main' into ananth7592-cosmos-native-aad-callback

f6049fb

ananth7592 marked this pull request as ready for review August 19, 2026 00:42

ananth7592 requested review from a team as code owners August 19, 2026 00:42

Copilot AI balanced review requested due to automatic review settings August 19, 2026 00:42

Copilot started reviewing on behalf of ananth7592 August 19, 2026 00:42 View session

Copilot AI reviewed Aug 19, 2026

View reviewed changes

Kevin Pilch (Pilchie) reviewed Aug 19, 2026

View reviewed changes

Tomas Varon (tvaron3) reviewed Aug 26, 2026

View reviewed changes

Navigation Menu

Uh oh!

Add native AAD credential callback - #5082

Add native AAD credential callback#5082
ananth7592 wants to merge 4 commits into
mainfrom
ananth7592-cosmos-native-aad-callback

ananth7592 commented Aug 18, 2026 •

edited

Loading

Uh oh!

azure-pipelines Bot commented Aug 18, 2026

Uh oh!

azure-pipelines Bot commented Aug 19, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Kevin Pilch (Pilchie) Aug 19, 2026

Uh oh!

Kevin Pilch (Pilchie) Aug 19, 2026

Uh oh!

Tomas Varon (tvaron3) Aug 26, 2026

Uh oh!

Tomas Varon (tvaron3) Aug 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		let completion_context =
		Box::into_raw(Box::new(PendingTokenRequest { sender })).cast::<c_void>();

Uh oh!

Conversation

ananth7592 commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add native AAD credential callback

Summary

Design

Initial token acquisition

Proactive refresh near expiry

Go v2 integration

Validation

Uh oh!

azure-pipelines Bot commented Aug 18, 2026

Uh oh!

azure-pipelines Bot commented Aug 19, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Kevin Pilch (Pilchie) Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Pilch (Pilchie) Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Tomas Varon (tvaron3) Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

Tomas Varon (tvaron3) Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ananth7592 commented Aug 18, 2026 •

edited

Loading