diff --git a/Cargo.lock b/Cargo.lock index fe9695bfd9c..2dff88b45a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -664,6 +664,7 @@ dependencies = [ name = "azure_data_cosmos_driver_native" version = "0.1.0" dependencies = [ + "async-trait", "azure_core 1.1.0", "azure_data_cosmos_driver", "cbindgen", diff --git a/sdk/cosmos/azure_data_cosmos_driver/docs/NATIVE_WRAPPER_SPEC.md b/sdk/cosmos/azure_data_cosmos_driver/docs/NATIVE_WRAPPER_SPEC.md index bc2b10f6e6d..edcb6c46a32 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/docs/NATIVE_WRAPPER_SPEC.md +++ b/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 diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml b/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml index 7d503a33d54..1f4225050bd 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/Cargo.toml +++ b/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 diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/build.rs b/sdk/cosmos/azure_data_cosmos_driver_native/build.rs index 7d46aa27509..7839791cd42 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/build.rs +++ b/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()), diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h index 742a156ba63..180140099ba 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h +++ b/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. */ diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/account_ref.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/account_ref.rs index ef874959cae..ff36c1c927e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/account_ref.rs +++ b/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` (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` 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)); + } + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/credential.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/credential.rs new file mode 100644 index 00000000000..11bb88abb13 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/credential.rs @@ -0,0 +1,598 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Host-provided token credential support for the native ABI. + +use std::{ffi::c_void, fmt, sync::Arc}; + +use async_trait::async_trait; +use azure_core::{ + credentials::{AccessToken, Secret, TokenCredential, TokenRequestOptions}, + error::ErrorKind, + time::{Duration, OffsetDateTime}, + Error, +}; +use tokio::sync::{oneshot, RwLock}; + +// Matches `azure_identity::TokenCache::should_refresh` in +// 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); + +/// 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. +pub type CosmosTokenCompletion = unsafe extern "C" fn( + completion_context: *mut c_void, + status: i32, + token: *const u8, + token_len: usize, + expires_on_unix_seconds: i64, + error_message: *const u8, + error_message_len: usize, +); + +/// Starts asynchronous token acquisition in the host. +/// +/// Returning zero transfers ownership of `request.completion_context` to the +/// host, which must eventually invoke `request.completion`. Returning nonzero +/// means the request was rejected synchronously and the completion must not be +/// invoked. +pub type CosmosTokenProviderCallback = + unsafe extern "C" fn(user_data: isize, request: *const CosmosTokenRequest) -> i32; + +/// Releases the host-owned state associated with a token provider. +pub type CosmosTokenProviderFree = unsafe extern "C" fn(user_data: isize); + +/// 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. +#[repr(C)] +pub struct CosmosTokenRequest { + /// UTF-8 token scope bytes. + pub scope: *const u8, + /// Number of bytes addressable from `scope`. + pub scope_len: usize, + /// Rust completion callback the host invokes exactly once. + pub completion: CosmosTokenCompletion, + /// Opaque Rust-owned context passed unchanged to `completion`. + pub completion_context: *mut c_void, +} + +/// Host callbacks used to acquire tokens and release host state. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct CosmosTokenProvider { + /// Starts token acquisition. Must be non-NULL. + pub get_token: + Option i32>, + /// Releases `user_data` after the last Rust credential reference is gone. + pub user_data_free: Option, +} + +struct PendingTokenRequest { + sender: oneshot::Sender>, +} + +struct CallbackTokenCredential { + provider: CosmosTokenProvider, + user_data: isize, + cached_token: RwLock>, +} + +impl CallbackTokenCredential { + fn new(provider: CosmosTokenProvider, user_data: isize) -> Option> { + provider.get_token?; + Some(Arc::new(Self { + provider, + user_data, + cached_token: RwLock::new(None), + })) + } + + async fn request_token(&self, scope: &str) -> azure_core::Result { + let (sender, receiver) = oneshot::channel(); + let completion_context = + Box::into_raw(Box::new(PendingTokenRequest { sender })).cast::(); + let status = { + let request = CosmosTokenRequest { + scope: scope.as_ptr(), + scope_len: scope.len(), + completion: complete_token_request, + completion_context, + }; + // SAFETY: the provider callback is supplied by the host under the C + // ABI contract. The request remains valid for the duration of the call. + unsafe { + (self + .provider + .get_token + .expect("validated when the credential was constructed"))( + self.user_data, + &request, + ) + } + }; + if status != 0 { + // The callback rejected the request synchronously, so ownership of + // the completion context did not transfer to the host. + // SAFETY: this pointer was allocated immediately above and the + // nonzero-return contract forbids the host from completing it. + unsafe { + drop(Box::from_raw( + completion_context.cast::(), + )); + } + return Err(credential_error(format!( + "host token provider rejected the request with status {status}" + ))); + } + + receiver.await.map_err(|_| { + credential_error("host token provider dropped the token request without completing it") + })? + } +} + +impl fmt::Debug for CallbackTokenCredential { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CallbackTokenCredential") + .field("provider", &"") + .field("user_data", &"") + .finish_non_exhaustive() + } +} + +impl Drop for CallbackTokenCredential { + fn drop(&mut self) { + if let Some(free) = self.provider.user_data_free { + // SAFETY: `user_data` is the opaque value supplied with this + // provider, and the callback is invoked exactly once on final drop. + unsafe { + free(self.user_data); + } + } + } +} + +#[async_trait] +impl TokenCredential for CallbackTokenCredential { + async fn get_token( + &self, + scopes: &[&str], + _options: Option>, + ) -> azure_core::Result { + let [scope] = scopes else { + return Err(credential_error( + "host token credentials require exactly one scope", + )); + }; + + // Fast path: shared read lock lets concurrent callers hit a warm + // cache in parallel. Mirrors `azure_identity::TokenCache` in + // sdk/identity/azure_identity/src/cache.rs. + { + let cached = self.cached_token.read().await; + let now = OffsetDateTime::now_utc(); + if let Some(token) = cached.as_ref() { + if token.expires_on > now.saturating_add(TOKEN_REFRESH_SKEW) { + return Ok(token.clone()); + } + } + } + + // 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; + let now = OffsetDateTime::now_utc(); + if let Some(token) = cached.as_ref() { + if token.expires_on > now.saturating_add(TOKEN_REFRESH_SKEW) { + return Ok(token.clone()); + } + } + + // Cache and return whatever the host handed us, matching + // `azure_identity::TokenCache::get_token` (which performs no + // post-fetch validation). If the host returns a token whose lifetime + // falls inside `TOKEN_REFRESH_SKEW`, the next call will see the + // fast-path check miss and re-invoke the host — same behavior a + // customer would see from any other Rust SDK credential in the same + // situation. Enforcing lifetime here would make this credential + // stricter than the rest of the SDK ecosystem. + let token = self.request_token(scope).await?; + *cached = Some(token.clone()); + Ok(token) + } +} + +pub(crate) fn create_token_credential( + provider: CosmosTokenProvider, + user_data: isize, +) -> Option> { + CallbackTokenCredential::new(provider, user_data) +} + +fn credential_error(message: impl Into) -> Error { + Error::with_message(ErrorKind::Credential, message.into()) +} + +unsafe extern "C" fn complete_token_request( + completion_context: *mut c_void, + status: i32, + token: *const u8, + token_len: usize, + expires_on_unix_seconds: i64, + error_message: *const u8, + error_message_len: usize, +) { + if completion_context.is_null() { + return; + } + + // SAFETY: ownership was transferred to the host only after the provider + // callback returned success. The host contract requires exactly one call. + let pending = unsafe { Box::from_raw(completion_context.cast::()) }; + let result = build_access_token( + status, + token, + token_len, + expires_on_unix_seconds, + error_message, + error_message_len, + ); + let _ = pending.sender.send(result); +} + +fn build_access_token( + status: i32, + token: *const u8, + token_len: usize, + expires_on_unix_seconds: i64, + error_message: *const u8, + error_message_len: usize, +) -> azure_core::Result { + if status != 0 { + // Mirror the sync-reject path (which includes the numeric status) so + // callers can classify AAD failures without string-matching against the + // host-supplied message. + let message = copy_host_bytes(error_message, error_message_len) + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|_| "no error message".to_string()); + return Err(credential_error(format!( + "host token provider failed with status {status}: {message}" + ))); + } + + let token = copy_host_bytes(token, token_len)?; + let token = String::from_utf8(token) + .map_err(|_| credential_error("host token provider returned a non-UTF-8 access token"))?; + if token.is_empty() { + return Err(credential_error( + "host token provider returned an empty access token", + )); + } + let expires_on = + OffsetDateTime::from_unix_timestamp(expires_on_unix_seconds).map_err(|_| { + credential_error("host token provider returned an invalid expiry timestamp") + })?; + Ok(AccessToken::new(Secret::new(token), expires_on)) +} + +fn copy_host_bytes(ptr: *const u8, len: usize) -> azure_core::Result> { + if len == 0 { + return Ok(Vec::new()); + } + if ptr.is_null() { + return Err(credential_error( + "host token provider returned a NULL buffer with a nonzero length", + )); + } + // SAFETY: the host guarantees the buffer is readable for `len` bytes for + // the duration of the completion callback. Copy before returning. + Ok(unsafe { std::slice::from_raw_parts(ptr, len) }.to_vec()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Mutex as StdMutex, + }; + + struct TestProvider { + calls: AtomicUsize, + frees: AtomicUsize, + tokens: StdMutex>, + scopes: StdMutex>, + } + + unsafe extern "C" fn get_token(user_data: isize, request: *const CosmosTokenRequest) -> i32 { + // SAFETY: tests pass a live `Arc` raw pointer as user data. + let state = unsafe { &*(user_data as *const TestProvider) }; + // SAFETY: the native adapter passes a valid request for this call. + let request = unsafe { &*request }; + // SAFETY: scope bytes are valid for the duration of this callback. + let scope = unsafe { std::slice::from_raw_parts(request.scope, request.scope_len) }; + state + .scopes + .lock() + .unwrap() + .push(String::from_utf8(scope.to_vec()).unwrap()); + state.calls.fetch_add(1, Ordering::SeqCst); + let (token, expires_on) = state.tokens.lock().unwrap().remove(0); + // SAFETY: token bytes remain alive until this synchronous completion + // call returns, and the completion copies them. + unsafe { + (request.completion)( + request.completion_context, + 0, + token.as_ptr(), + token.len(), + expires_on, + std::ptr::null(), + 0, + ); + } + 0 + } + + unsafe extern "C" fn get_token_async( + user_data: isize, + request: *const CosmosTokenRequest, + ) -> i32 { + // SAFETY: tests pass a live `Arc` raw pointer as user data. + let state = unsafe { &*(user_data as *const TestProvider) }; + // SAFETY: the native adapter passes a valid request for this call. + let request = unsafe { &*request }; + // SAFETY: scope bytes are valid for the duration of this callback. + let scope = unsafe { std::slice::from_raw_parts(request.scope, request.scope_len) }; + state + .scopes + .lock() + .unwrap() + .push(String::from_utf8(scope.to_vec()).unwrap()); + state.calls.fetch_add(1, Ordering::SeqCst); + let (token, expires_on) = state.tokens.lock().unwrap().remove(0); + let completion = request.completion; + let completion_context = request.completion_context as usize; + std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(20)); + // SAFETY: the provider returned success, transferring this context + // to the host. Token bytes remain live until completion returns. + unsafe { + completion( + completion_context as *mut c_void, + 0, + token.as_ptr(), + token.len(), + expires_on, + std::ptr::null(), + 0, + ); + } + }); + 0 + } + + unsafe extern "C" fn reject_token( + _user_data: isize, + _request: *const CosmosTokenRequest, + ) -> i32 { + 17 + } + + unsafe extern "C" fn fail_token(_user_data: isize, request: *const CosmosTokenRequest) -> i32 { + const MESSAGE: &[u8] = b"credential unavailable"; + // SAFETY: the native adapter passes a valid request for this call. + let request = unsafe { &*request }; + // SAFETY: the message remains live for the duration of completion. + unsafe { + (request.completion)( + request.completion_context, + 1, + std::ptr::null(), + 0, + 0, + MESSAGE.as_ptr(), + MESSAGE.len(), + ); + } + 0 + } + + unsafe extern "C" fn free_provider(user_data: isize) { + // SAFETY: ownership of this strong reference was transferred to the + // credential at construction. + let state = unsafe { Arc::from_raw(user_data as *const TestProvider) }; + state.frees.fetch_add(1, Ordering::SeqCst); + } + + fn credential(tokens: Vec<(String, i64)>) -> (Arc, Arc) { + credential_with_callback(tokens, get_token) + } + + fn credential_with_callback( + tokens: Vec<(String, i64)>, + get_token: CosmosTokenProviderCallback, + ) -> (Arc, Arc) { + let state = Arc::new(TestProvider { + calls: AtomicUsize::new(0), + frees: AtomicUsize::new(0), + tokens: StdMutex::new(tokens), + scopes: StdMutex::new(Vec::new()), + }); + let user_data = Arc::into_raw(Arc::clone(&state)) as isize; + let credential = create_token_credential( + CosmosTokenProvider { + get_token: Some(get_token), + user_data_free: Some(free_provider), + }, + user_data, + ) + .unwrap(); + (state, credential) + } + + #[tokio::test] + async fn caches_token_until_refresh_window() { + let expires = OffsetDateTime::now_utc() + .saturating_add(Duration::minutes(10)) + .unix_timestamp(); + let (state, credential) = credential(vec![("token-a".to_string(), expires)]); + + let first = credential.get_token(&["scope"], None).await.unwrap(); + let second = credential.get_token(&["scope"], None).await.unwrap(); + + assert_eq!(first.token.secret(), "token-a"); + assert_eq!(second.token.secret(), "token-a"); + assert_eq!(state.calls.load(Ordering::SeqCst), 1); + assert_eq!(state.scopes.lock().unwrap().as_slice(), ["scope"]); + } + + #[tokio::test] + async fn refreshes_token_near_expiry() { + let now = OffsetDateTime::now_utc(); + let (state, credential) = credential(vec![ + ( + "token-a".to_string(), + now.saturating_add(Duration::seconds(30)).unix_timestamp(), + ), + ( + "token-b".to_string(), + now.saturating_add(Duration::minutes(10)).unix_timestamp(), + ), + ]); + + let first = credential.get_token(&["scope"], None).await.unwrap(); + let second = credential.get_token(&["scope"], None).await.unwrap(); + + assert_eq!(first.token.secret(), "token-a"); + assert_eq!(second.token.secret(), "token-b"); + assert_eq!(state.calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn awaits_asynchronous_host_completion() { + let expires = OffsetDateTime::now_utc() + .saturating_add(Duration::minutes(10)) + .unix_timestamp(); + let (state, credential) = + credential_with_callback(vec![("token-a".to_string(), expires)], get_token_async); + + let token = credential.get_token(&["scope"], None).await.unwrap(); + + assert_eq!(token.token.secret(), "token-a"); + assert_eq!(state.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn concurrent_requests_single_flight_refresh() { + let expires = OffsetDateTime::now_utc() + .saturating_add(Duration::minutes(10)) + .unix_timestamp(); + let (state, credential) = + credential_with_callback(vec![("token-a".to_string(), expires)], get_token_async); + + let (a, b, c, d) = tokio::join!( + credential.get_token(&["scope"], None), + credential.get_token(&["scope"], None), + credential.get_token(&["scope"], None), + credential.get_token(&["scope"], None), + ); + + for token in [a, b, c, d] { + assert_eq!(token.unwrap().token.secret(), "token-a"); + } + assert_eq!(state.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn propagates_synchronous_provider_rejection() { + let (_state, credential) = credential_with_callback(Vec::new(), reject_token); + + let error = credential.get_token(&["scope"], None).await.unwrap_err(); + + assert!(error.to_string().contains("status 17")); + } + + #[tokio::test] + async fn propagates_host_completion_error() { + let (_state, credential) = credential_with_callback(Vec::new(), fail_token); + + let error = credential.get_token(&["scope"], None).await.unwrap_err(); + + let message = error.to_string(); + assert!( + message.contains("status 1"), + "missing status code: {message}" + ); + assert!( + message.contains("credential unavailable"), + "missing host message: {message}" + ); + } + + #[tokio::test] + async fn returns_short_lived_host_token_verbatim() { + // Matches `azure_identity::TokenCache`, which performs no post-fetch + // validation of the token lifetime. A host returning a very short + // (or already-expired) token gets cached and returned as-is; the + // consumer will see refresh thrash on subsequent calls, same as any + // other Rust SDK credential would produce in the same situation. + let expires = OffsetDateTime::now_utc() + .saturating_sub(Duration::seconds(1)) + .unix_timestamp(); + let (state, credential) = credential(vec![("short-lived".to_string(), expires)]); + + let token = credential + .get_token(&["scope"], None) + .await + .expect("short-lived tokens should be returned verbatim"); + + assert_eq!(token.token.secret(), "short-lived"); + assert_eq!(state.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn rejects_multiple_scopes_without_calling_host() { + let (state, credential) = credential(Vec::new()); + + let error = credential + .get_token(&["scope-a", "scope-b"], None) + .await + .unwrap_err(); + + assert!(error.to_string().contains("exactly one scope")); + assert_eq!(state.calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn frees_host_state_after_final_credential_drop() { + let expires = OffsetDateTime::now_utc() + .saturating_add(Duration::minutes(10)) + .unix_timestamp(); + let (state, credential) = credential(vec![("token-a".to_string(), expires)]); + + assert_eq!(state.frees.load(Ordering::SeqCst), 0); + drop(credential); + assert_eq!(state.frees.load(Ordering::SeqCst), 1); + } + + #[test] + fn rejects_missing_provider_callback() { + let credential = create_token_credential( + CosmosTokenProvider { + get_token: None, + user_data_free: None, + }, + 0, + ); + assert!(credential.is_none()); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs index 8ef30c16d5a..b0899b2cca7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/lib.rs @@ -25,6 +25,7 @@ pub mod account_ref; pub mod bytes; pub mod completion; pub mod container_ref; +pub mod credential; pub mod database_ref; pub mod driver; pub mod driver_options;