From 37db397034d1ea2e84aa7991da9c19c916f38273 Mon Sep 17 00:00:00 2001 From: John Batty Date: Thu, 17 Aug 2023 11:36:06 +0100 Subject: [PATCH] Implement Cosmos AAD authentication --- sdk/data_cosmos/src/authorization_policy.rs | 51 ++++++++++++++----- .../permission/authorization_token.rs | 27 +++++++++- .../resources/permission/permission_token.rs | 5 ++ 3 files changed, 68 insertions(+), 15 deletions(-) diff --git a/sdk/data_cosmos/src/authorization_policy.rs b/sdk/data_cosmos/src/authorization_policy.rs index 1c8698de37..ed11f934cd 100644 --- a/sdk/data_cosmos/src/authorization_policy.rs +++ b/sdk/data_cosmos/src/authorization_policy.rs @@ -26,7 +26,7 @@ const VERSION_NUMBER: &str = "1.0"; /// /// This struct implements `Debug` but secrets are encrypted by `AuthorizationToken` so there is no risk of /// leaks in debug logs (secrets are stored in cleartext in memory: dumps are still leaky). -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone)] pub struct AuthorizationPolicy { authorization_token: AuthorizationToken, } @@ -67,6 +67,7 @@ impl Policy for AuthorizationPolicy { &resource_link, time_nonce, ) + .await? }; trace!( @@ -136,17 +137,23 @@ fn generate_resource_link(request: &Request) -> String { } } -/// The `CosmosDB` authorization can either be "primary" (i.e., one of the two service-level tokens) or -/// "resource" (i.e., a single database). In the first case the signature must be constructed by -/// signing the HTTP method, resource type, resource link (the relative URI) and the current time. -/// In the second case, the signature is just the resource key. -fn generate_authorization( +/// The `CosmosDB` authorization can either be: +/// - "primary": one of the two service-level tokens +/// - "resource: e.g. a single database +/// - "aad": Azure Active Directory token +/// In the "primary" case the signature must be constructed by signing the HTTP method, +/// resource type, resource link (the relative URI) and the current time. +/// +/// In the "resource" case, the signature is just the resource key. +/// +/// In the "aad" case, the signature is the AAD token. +async fn generate_authorization( auth_token: &AuthorizationToken, http_method: &azure_core::Method, resource_type: &ResourceType, resource_link: &str, time_nonce: OffsetDateTime, -) -> String { +) -> azure_core::Result { let (authorization_type, signature) = match auth_token { AuthorizationToken::Primary(key) => { let string_to_sign = @@ -157,6 +164,17 @@ fn generate_authorization( ) } AuthorizationToken::Resource(key) => ("resource", Cow::Borrowed(key)), + AuthorizationToken::TokenCredential(token_credential) => ( + "aad", + Cow::Owned( + token_credential + .get_token(resource_link) + .await? + .token + .secret() + .to_string(), + ), + ), }; let str_unencoded = format!("type={authorization_type}&ver={VERSION_NUMBER}&sig={signature}"); @@ -165,7 +183,7 @@ fn generate_authorization( str_unencoded ); - form_urlencoded::byte_serialize(str_unencoded.as_bytes()).collect::() + Ok(form_urlencoded::byte_serialize(str_unencoded.as_bytes()).collect::()) } /// This function generates a valid authorization string, according to the documentation. @@ -255,8 +273,8 @@ mon, 01 jan 1900 01:00:00 gmt ); } - #[test] - fn generate_authorization_00() { + #[tokio::test] + async fn generate_authorization_00() { let time = date::parse_rfc3339("1900-01-01T01:00:00.000000000+00:00").unwrap(); let auth_token = AuthorizationToken::primary_from_base64( @@ -270,15 +288,18 @@ mon, 01 jan 1900 01:00:00 gmt &ResourceType::Databases, "dbs/MyDatabase/colls/MyCollection", time, - ); + ) + .await + .unwrap(); + assert_eq!( ret, "type%3Dmaster%26ver%3D1.0%26sig%3DQkz%2Fr%2B1N2%2BPEnNijxGbGB%2FADvLsLBQmZ7uBBMuIwf4I%3D" ); } - #[test] - fn generate_authorization_01() { + #[tokio::test] + async fn generate_authorization_01() { let time = date::parse_rfc3339("2017-04-27T00:51:12.000000000+00:00").unwrap(); let auth_token = AuthorizationToken::primary_from_base64( @@ -292,7 +313,9 @@ mon, 01 jan 1900 01:00:00 gmt &ResourceType::Databases, "dbs/ToDoList", time, - ); + ) + .await + .unwrap(); // This is the result shown in the MSDN page. It's clearly wrong :) // below is the correct one. diff --git a/sdk/data_cosmos/src/resources/permission/authorization_token.rs b/sdk/data_cosmos/src/resources/permission/authorization_token.rs index af15fa6486..cb5feb52f6 100644 --- a/sdk/data_cosmos/src/resources/permission/authorization_token.rs +++ b/sdk/data_cosmos/src/resources/permission/authorization_token.rs @@ -1,21 +1,40 @@ use super::PermissionToken; +use azure_core::auth::TokenCredential; use azure_core::{ base64, error::{Error, ErrorKind}, }; use std::fmt; +use std::sync::Arc; /// Authorization tokens for accessing Cosmos. /// /// Learn more about the different types of tokens [here](https://docs.microsoft.com/azure/cosmos-db/secure-access-to-data). -#[derive(PartialEq, Clone, Eq)] +#[derive(Clone)] pub enum AuthorizationToken { /// Used for administrative resources: database accounts, databases, users, and permissions Primary(Vec), /// Used for application resources: containers, documents, attachments, stored procedures, triggers, and UDFs Resource(String), + /// AAD token credential + TokenCredential(Arc), } +impl PartialEq for AuthorizationToken { + fn eq(&self, other: &Self) -> bool { + use AuthorizationToken::*; + match (self, other) { + (Primary(a), Primary(b)) => a == b, + (Resource(a), Resource(b)) => a == b, + // Consider two token credentials equal if they point to the same object. + (TokenCredential(a), TokenCredential(b)) => Arc::ptr_eq(a, b), + _ => false, + } + } +} + +impl Eq for AuthorizationToken {} + impl AuthorizationToken { /// Create a primary `AuthorizationToken` from base64 encoded data /// @@ -32,6 +51,11 @@ impl AuthorizationToken { pub fn new_resource(resource: String) -> AuthorizationToken { AuthorizationToken::Resource(resource) } + + /// Create an `AuthorizationToken` from a `TokenCredential`. + pub fn from_token_credential(token_credential: Arc) -> AuthorizationToken { + AuthorizationToken::TokenCredential(token_credential) + } } impl fmt::Debug for AuthorizationToken { @@ -43,6 +67,7 @@ impl fmt::Debug for AuthorizationToken { match self { AuthorizationToken::Primary(_) => "Master", AuthorizationToken::Resource(_) => "Resource", + AuthorizationToken::TokenCredential(_) => "TokenCredential", } ) } diff --git a/sdk/data_cosmos/src/resources/permission/permission_token.rs b/sdk/data_cosmos/src/resources/permission/permission_token.rs index 9454f3e26e..d1aa88371a 100644 --- a/sdk/data_cosmos/src/resources/permission/permission_token.rs +++ b/sdk/data_cosmos/src/resources/permission/permission_token.rs @@ -30,6 +30,11 @@ impl std::fmt::Display for PermissionToken { let (permission_type, signature) = match &self.token { AuthorizationToken::Resource(s) => ("resource", Cow::Borrowed(s)), AuthorizationToken::Primary(s) => ("master", Cow::Owned(base64::encode(s))), + // @@@TODO: Do we need to support TokenCredential for PermissionToken? + // It is painful to implement because we can't easily get the + // token string from the TokenCredential (the function is async and fallible, + // and fmt() is neither!). + AuthorizationToken::TokenCredential(_s) => ("aad", Cow::Owned("xxx".to_string())), }; write!( f,