Add native AAD credential callback

sdk/cosmos/azure_data_cosmos_driver/docs/NATIVE_WRAPPER_SPEC.md

@@ -1092,9 +1092,8 @@ cosmos_status_code_t cosmos_account_ref_with_master_key(
  * its user_data are kept alive by the AccountReference's Arc. */
 cosmos_status_code_t cosmos_account_ref_with_credential(
     const char *endpoint,
-    cosmos_token_provider_t credential,   /* see §4.10 */
-    void *user_data,
-    void (*user_data_free)(void *user_data),
+    cosmos_token_provider_t provider,   /* see §4.9 */
+    intptr_t user_data,
     cosmos_account_ref_t **out_account,
     cosmos_error_t **out_error);
 
@@ -1632,6 +1631,49 @@ void cosmos_diagnostics_free(cosmos_diagnostics_t *d);
 
 The JSON snapshot is the **only** place the wrapper serializes anything to JSON, and it's purely a debugging aid — schema-agnosticism is preserved on the data plane.
 
+### 4.9 Host token credentials (`src/credential.rs`)
+
+AAD credentials remain owned by the host SDK. The native wrapper adapts an asynchronous host callback into `azure_core::credentials::TokenCredential`, caches the returned token, and single-flights refresh so concurrent Cosmos requests do not invoke the host credential in parallel.
+
+```c
+typedef void (*cosmos_token_completion_t)(
+    void *completion_context,
+    int32_t status,
+    const uint8_t *token,
+    uintptr_t token_len,
+    int64_t expires_on_unix_seconds,
+    const uint8_t *error_message,
+    uintptr_t error_message_len);
+
+typedef struct cosmos_token_request_t {
+    const uint8_t *scope;
+    uintptr_t scope_len;
+    cosmos_token_completion_t completion;
+    void *completion_context;
+} cosmos_token_request_t;
+
+typedef struct cosmos_token_provider_t {
+    int32_t (*get_token)(
+        intptr_t user_data,
+        const cosmos_token_request_t *request);
+    void (*user_data_free)(intptr_t user_data); /* nullable */
+} cosmos_token_provider_t;
+```
+
+`get_token` is required. `user_data_free` is optional. `user_data` is a pointer-sized opaque integer rather than a host pointer; Go bindings use it for a `runtime/cgo.Handle` value and must not retain a Go pointer in C or Rust memory.
+
+The request follows this ownership state machine:
+
+1. Rust owns `completion_context` while calling `get_token`. `scope` and the request struct are borrowed only for that call, so the host copies the scope before starting asynchronous work.
+2. A nonzero `get_token` return rejects synchronously. Rust retains and destroys `completion_context`; the host must not call `completion`.
+3. A zero return transfers `completion_context` to the host. The host must call `completion` exactly once, even when token acquisition fails.
+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.
+
+The callback may complete synchronously before `get_token` returns or asynchronously from another host thread. Go must invoke the completion function through a small C trampoline because cgo cannot call a C function pointer directly. Cancellation of a Cosmos operation does not cancel an already accepted token request in v1: after returning zero, the host still owns the completion context and must complete it. A future ABI revision can add an explicit cancellation hook without changing token or account ownership.
+
 ---
 
 ## 5. Build & Distribution

sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml

@@ -32,6 +32,7 @@ azure_data_cosmos_driver = { path = "../azure_data_cosmos_driver", version = "0.
 # endpoint. Both are zero-cost wrappers around existing driver deps so
 # they don't grow the closure.
 azure_core = { workspace = true, default-features = false }
+async-trait.workspace = true
 url = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "sync", "time", "macros"] }
 # `FutureExt::catch_unwind` is used by the submit pipeline to keep the

sdk/cosmos/azure_data_cosmos_driver_native/build.rs

@@ -101,6 +101,9 @@ fn generate_c_header() {
         ("DriverHandle".into(), "driver_t".into()),
         ("AccountReference".into(), "account_ref_t".into()),
         ("AccountRefHandle".into(), "account_ref_t".into()),
+        ("CosmosTokenCompletion".into(), "token_completion_t".into()),
+        ("CosmosTokenRequest".into(), "token_request_t".into()),
+        ("CosmosTokenProvider".into(), "token_provider_t".into()),
         ("DatabaseReference".into(), "database_ref_t".into()),
         ("DatabaseRefHandle".into(), "database_ref_t".into()),
         ("ContainerReference".into(), "container_ref_t".into()),

sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h

@@ -961,6 +961,60 @@ typedef struct cosmos_error_t {
   const char *backtrace;
 } cosmos_error_t;
 
+/**
+ * Completes an asynchronous host token request.
+ *
+ * The host must invoke this callback exactly once after its token-provider
+ * callback returns success. Token and error buffers are borrowed only for the
+ * duration of this call; Rust copies them before returning.
+ */
+typedef void (*cosmos_token_completion_t)(void *completion_context,
+                                          int32_t status,
+                                          const uint8_t *token,
+                                          uintptr_t token_len,
+                                          int64_t expires_on_unix_seconds,
+                                          const uint8_t *error_message,
+                                          uintptr_t error_message_len);
+
+/**
+ * One asynchronous access-token request passed to the host.
+ *
+ * `scope` is borrowed and valid only until the token-provider callback
+ * returns. The host must copy it before starting asynchronous work.
+ */
+typedef struct cosmos_token_request_t {
+  /**
+   * UTF-8 token scope bytes.
+   */
+  const uint8_t *scope;
+  /**
+   * Number of bytes addressable from `scope`.
+   */
+  uintptr_t scope_len;
+  /**
+   * Rust completion callback the host invokes exactly once.
+   */
+  cosmos_token_completion_t completion;
+  /**
+   * Opaque Rust-owned context passed unchanged to `completion`.
+   */
+  void *completion_context;
+} cosmos_token_request_t;
+
+/**
+ * Host callbacks used to acquire tokens and release host state.
+ */
+typedef struct cosmos_token_provider_t {
+  /**
+   * Starts token acquisition. Must be non-NULL.
+   */
+  int32_t (*get_token)(intptr_t user_data, const struct cosmos_token_request_t *request);
+  /**
+   * Releases `user_data` after the last Rust credential reference is gone.
+   */
+  void (*user_data_free)(intptr_t user_data);
+} cosmos_token_provider_t;
+
 /**
  * A library-owned byte buffer returned by value across the C ABI.
  *
@@ -1644,6 +1698,29 @@ cosmos_status_code_t cosmos_account_ref_with_master_key(const char *endpoint,
                                                         struct cosmos_account_ref_t **out_account,
                                                         struct cosmos_error_t **out_error);
 
+/**
+ * Creates an account reference authenticated by a host token credential.
+ *
+ * The callback provider is adapted into the driver's async
+ * [`azure_core::credentials::TokenCredential`] interface. Ownership of
+ * `user_data` transfers to Rust only on success. The optional
+ * `user_data_free` callback runs after the final account/driver credential
+ * reference is released.
+ *
+ * # Returns
+ *
+ * - `SUCCESS` (0) with `*out_account` populated.
+ * - `INVALID_ARGUMENT` (1) when `endpoint`, `out_account`, or the provider's
+ *   `get_token` callback is NULL.
+ * - `INVALID_UTF8` (2) when `endpoint` is not valid UTF-8.
+ * - `INVALID_ACCOUNT_REFERENCE` (4003) when `endpoint` is not a parsable URL.
+ */
+cosmos_status_code_t cosmos_account_ref_with_credential(const char *endpoint,
+                                                        struct cosmos_token_provider_t provider,
+                                                        intptr_t user_data,
+                                                        struct cosmos_account_ref_t **out_account,
+                                                        struct cosmos_error_t **out_error);
+
 /**
  * Frees an account-reference handle. NULL is a no-op.
  */

sdk/cosmos/azure_data_cosmos_driver_native/src/account_ref.rs

@@ -4,12 +4,8 @@
 //! C ABI surface for `cosmos_account_ref_t` — wraps the driver's
 //! [`azure_data_cosmos_driver::models::AccountReference`].
 //!
-//! The wrapper currently ships only the master-key path. Token-credential
-//! (`AccountReference::with_credential`) and resource-token paths require an
-//! FFI bridge for `Arc<dyn TokenCredential>` (an async trait whose
-//! implementations live in `azure_identity`) — bridging an arbitrary C-side
-//! async credential through FFI is non-trivial and is intentionally
-//! deferred to a follow-up.
+//! Master-key credentials are copied into Rust-owned memory. Token credentials
+//! use a host callback adapted to the driver's async `TokenCredential` trait.
 //!
 //! Construction validates the endpoint URL up-front; a parse failure
 //! surfaces a `400 Bad Request` packed status whose sub-status is
@@ -27,7 +23,10 @@ use azure_core::credentials::Secret;
 use azure_data_cosmos_driver::models::AccountReference as DriverAccountReference;
 use url::Url;
 
-use crate::error::{CosmosError, CosmosErrorCode, CosmosStatusCode};
+use crate::{
+    credential::{create_token_credential, CosmosTokenProvider},
+    error::{CosmosError, CosmosErrorCode, CosmosStatusCode},
+};
 
 /// The C ABI handle for an account reference (`cosmos_account_ref_t`).
 ///
@@ -195,6 +194,51 @@ pub extern "C" fn cosmos_account_ref_with_master_key(
     CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code()
 }
 
+/// Creates an account reference authenticated by a host token credential.
+///
+/// The callback provider is adapted into the driver's async
+/// [`azure_core::credentials::TokenCredential`] interface. Ownership of
+/// `user_data` transfers to Rust only on success. The optional
+/// `user_data_free` callback runs after the final account/driver credential
+/// reference is released.
+///
+/// # Returns
+///
+/// - `SUCCESS` (0) with `*out_account` populated.
+/// - `INVALID_ARGUMENT` (1) when `endpoint`, `out_account`, or the provider's
+///   `get_token` callback is NULL.
+/// - `INVALID_UTF8` (2) when `endpoint` is not valid UTF-8.
+/// - `INVALID_ACCOUNT_REFERENCE` (4003) when `endpoint` is not a parsable URL.
+#[no_mangle]
+pub extern "C" fn cosmos_account_ref_with_credential(
+    endpoint: *const c_char,
+    provider: CosmosTokenProvider,
+    user_data: isize,
+    out_account: *mut *mut AccountRefHandle,
+    out_error: *mut *mut CosmosError,
+) -> CosmosStatusCode {
+    if out_account.is_null() || provider.get_token.is_none() {
+        return CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code();
+    }
+    let endpoint_str = match try_cstr_to_str(endpoint) {
+        Ok(s) => s,
+        Err(code) => return code.as_status_code(),
+    };
+    let url = match parse_endpoint(endpoint_str, out_error) {
+        Ok(u) => u,
+        Err(code) => return code.as_status_code(),
+    };
+    let credential = create_token_credential(provider, user_data)
+        .expect("provider callback was validated before credential construction");
+    let driver_ref = DriverAccountReference::with_credential(url, credential);
+    let handle = AccountRefHandle::into_raw(driver_ref);
+    // SAFETY: caller guarantees `out_account` is writable for one handle.
+    unsafe {
+        *out_account = handle;
+    }
+    CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code()
+}
+
 /// Frees an account-reference handle. NULL is a no-op.
 #[no_mangle]
 pub extern "C" fn cosmos_account_ref_free(account: *mut AccountRefHandle) {
@@ -210,6 +254,10 @@ pub(crate) mod tests {
     use super::*;
     use std::ffi::CString;
     use std::ptr;
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
 
     fn ok_cstr(s: &str) -> CString {
         CString::new(s).expect("test inputs must be NUL-free")
@@ -297,4 +345,103 @@ pub(crate) mod tests {
         );
         assert!(out.is_null());
     }
+
+    unsafe extern "C" fn unused_get_token(
+        _user_data: isize,
+        _request: *const crate::credential::CosmosTokenRequest,
+    ) -> i32 {
+        unreachable!("credential callback is not used during construction")
+    }
+
+    unsafe extern "C" fn free_provider(user_data: isize) {
+        // SAFETY: tests transfer one strong `Arc<AtomicUsize>` reference.
+        let released = unsafe { Arc::from_raw(user_data as *const AtomicUsize) };
+        released.fetch_add(1, Ordering::SeqCst);
+    }
+
+    #[test]
+    fn with_credential_rejects_null_out_account_without_taking_ownership() {
+        let endpoint = ok_cstr("https://x.documents.azure.com:443/");
+        let released = Arc::new(AtomicUsize::new(0));
+        let user_data = Arc::into_raw(Arc::clone(&released)) as isize;
+        let provider = CosmosTokenProvider {
+            get_token: Some(unused_get_token),
+            user_data_free: Some(free_provider),
+        };
+
+        let rc = cosmos_account_ref_with_credential(
+            endpoint.as_ptr(),
+            provider,
+            user_data,
+            ptr::null_mut(),
+            ptr::null_mut(),
+        );
+
+        assert_eq!(
+            rc,
+            CosmosErrorCode::CosmosErrorCodeInvalidArgument.as_status_code()
+        );
+        assert_eq!(released.load(Ordering::SeqCst), 0);
+        // SAFETY: ownership was not transferred, so reclaim the extra strong
+        // reference created for the attempted constructor.
+        unsafe {
+            drop(Arc::from_raw(user_data as *const AtomicUsize));
+        }
+    }
+
+    #[test]
+    fn with_credential_releases_host_state_with_account() {
+        let endpoint = ok_cstr("https://x.documents.azure.com:443/");
+        let released = Arc::new(AtomicUsize::new(0));
+        let user_data = Arc::into_raw(Arc::clone(&released)) as isize;
+        let provider = CosmosTokenProvider {
+            get_token: Some(unused_get_token),
+            user_data_free: Some(free_provider),
+        };
+        let mut out: *mut AccountRefHandle = ptr::null_mut();
+
+        let rc = cosmos_account_ref_with_credential(
+            endpoint.as_ptr(),
+            provider,
+            user_data,
+            &mut out,
+            ptr::null_mut(),
+        );
+
+        assert_eq!(rc, CosmosErrorCode::CosmosErrorCodeSuccess.as_status_code());
+        assert!(!out.is_null());
+        cosmos_account_ref_free(out);
+        assert_eq!(released.load(Ordering::SeqCst), 1);
+    }
+
+    #[test]
+    fn with_credential_invalid_endpoint_does_not_take_ownership() {
+        let endpoint = ok_cstr("not a url");
+        let released = Arc::new(AtomicUsize::new(0));
+        let user_data = Arc::into_raw(Arc::clone(&released)) as isize;
+        let provider = CosmosTokenProvider {
+            get_token: Some(unused_get_token),
+            user_data_free: Some(free_provider),
+        };
+        let mut out: *mut AccountRefHandle = ptr::null_mut();
+
+        let rc = cosmos_account_ref_with_credential(
+            endpoint.as_ptr(),
+            provider,
+            user_data,
+            &mut out,
+            ptr::null_mut(),
+        );
+
+        assert_eq!(
+            rc,
+            CosmosErrorCode::CosmosErrorCodeInvalidAccountReference.as_status_code()
+        );
+        assert!(out.is_null());
+        assert_eq!(released.load(Ordering::SeqCst), 0);
+        // SAFETY: ownership was not transferred on constructor failure.
+        unsafe {
+            drop(Arc::from_raw(user_data as *const AtomicUsize));
+        }
+    }
 }