Implement Cosmos AAD authentication (#1338)

sdk/data_cosmos/src/authorization_policy.rs

@@ -2,7 +2,7 @@ use crate::resources::permission::AuthorizationToken;
 use crate::resources::ResourceType;
 use azure_core::base64;
 use azure_core::headers::{HeaderValue, AUTHORIZATION, MS_DATE, VERSION};
-use azure_core::{date, Context, Policy, PolicyResult, Request};
+use azure_core::{date, Context, Policy, PolicyResult, Request, Url};
 use hmac::{Hmac, Mac};
 use sha2::Sha256;
 use std::borrow::Cow;
@@ -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,
 }
@@ -62,11 +62,13 @@ impl Policy for AuthorizationPolicy {
             generate_authorization(
                 &self.authorization_token,
                 request.method(),
+                request.url(),
                 ctx.get()
                     .expect("ResourceType must be in the Context at this point"),
                 &resource_link,
                 time_nonce,
             )
+            .await?
         };
 
         trace!(
@@ -136,17 +138,24 @@ 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,
+    url: &Url,
     resource_type: &ResourceType,
     resource_link: &str,
     time_nonce: OffsetDateTime,
-) -> String {
+) -> azure_core::Result<String> {
     let (authorization_type, signature) = match auth_token {
         AuthorizationToken::Primary(key) => {
             let string_to_sign =
@@ -157,6 +166,17 @@ fn generate_authorization(
             )
         }
         AuthorizationToken::Resource(key) => ("resource", Cow::Borrowed(key)),
+        AuthorizationToken::TokenCredential(token_credential) => (
+            "aad",
+            Cow::Owned(
+                token_credential
+                    .get_token(&scope_from_url(url))
+                    .await?
+                    .token
+                    .secret()
+                    .to_string(),
+            ),
+        ),
     };
 
     let str_unencoded = format!("type={authorization_type}&ver={VERSION_NUMBER}&sig={signature}");
@@ -165,7 +185,15 @@ fn generate_authorization(
         str_unencoded
     );
 
-    form_urlencoded::byte_serialize(str_unencoded.as_bytes()).collect::<String>()
+    Ok(form_urlencoded::byte_serialize(str_unencoded.as_bytes()).collect::<String>())
+}
+
+/// This function generates the scope string from the passed url. The scope string is used to
+/// request the AAD token.
+fn scope_from_url(url: &Url) -> String {
+    let scheme = url.scheme();
+    let hostname = url.host_str().unwrap();
+    format!("{scheme}://{hostname}")
 }
 
 /// This function generates a valid authorization string, according to the documentation.
@@ -255,44 +283,55 @@ 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(
             "8F8xXXOptJxkblM1DBXW7a6NMI5oE8NnwPGYBmwxLCKfejOK7B7yhcCHMGvN3PBrlMLIOeol1Hv9RCdzAZR5sg==",
         )
         .unwrap();
 
+        let url = azure_core::Url::parse("https://.documents.azure.com/dbs/ToDoList").unwrap();
+
         let ret = generate_authorization(
             &auth_token,
             &azure_core::Method::Get,
+            &url,
             &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(
             "dsZQi3KtZmCv1ljt3VNWNm7sQUF1y5rJfC6kv5JiwvW0EndXdDku/dkKBp8/ufDToSxL",
         )
         .unwrap();
 
+        let url = azure_core::Url::parse("https://.documents.azure.com/dbs/ToDoList").unwrap();
+
         let ret = generate_authorization(
             &auth_token,
             &azure_core::Method::Get,
+            &url,
             &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.
@@ -340,4 +379,12 @@ mon, 01 jan 1900 01:00:00 gmt
         );
         assert_eq!(&generate_resource_link(&request), "dbs/test_db");
     }
+
+    #[test]
+    fn scope_from_url_01() {
+        let scope = scope_from_url(
+            &azure_core::Url::parse("https://.documents.azure.com/dbs/test_db/colls").unwrap(),
+        );
+        assert_eq!(scope, "https://.documents.azure.com");
+    }
 }

sdk/data_cosmos/src/resources/permission/authorization_token.rs

@@ -1,19 +1,23 @@
 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<u8>),
     /// Used for application resources: containers, documents, attachments, stored procedures, triggers, and UDFs
     Resource(String),
+    /// AAD token credential
+    TokenCredential(Arc<dyn TokenCredential>),
 }
 
 impl AuthorizationToken {
@@ -32,6 +36,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<dyn TokenCredential>) -> AuthorizationToken {
+        AuthorizationToken::TokenCredential(token_credential)
+    }
 }
 
 impl fmt::Debug for AuthorizationToken {
@@ -43,6 +52,7 @@ impl fmt::Debug for AuthorizationToken {
             match self {
                 AuthorizationToken::Primary(_) => "Master",
                 AuthorizationToken::Resource(_) => "Resource",
+                AuthorizationToken::TokenCredential(_) => "TokenCredential",
             }
         )
     }

sdk/data_cosmos/src/resources/permission/permission.rs

@@ -10,7 +10,8 @@ use std::borrow::Cow;
 /// access to a specific resource. It is used to manage access to collections, documents,
 /// attachments, stored procedures, triggers, and user-defined functions for a particular user.
 /// You can learn more about permissions [here](https://docs.microsoft.com/rest/api/cosmos-db/permissions).
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+//#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Permission {
     ///  The unique name that identifies the permission.
     pub id: String,

sdk/data_cosmos/src/resources/permission/permission_response.rs

@@ -5,7 +5,7 @@ use azure_core::Response as HttpResponse;
 
 use super::Permission;
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone)]
 pub struct PermissionResponse {
     pub permission: Permission,
     pub charge: f64,

sdk/data_cosmos/src/resources/permission/permission_token.rs

@@ -9,12 +9,25 @@ const SIGNATURE_PREFIX: &str = "sig=";
 ///
 /// This field is a url encoded string with the type of permission, the signature, and the version (currently only 1.0)
 /// This type is a wrapper around `AuthorizationToken`.
-#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
+#[derive(Debug, Clone, Deserialize)]
 #[serde(try_from = "String")]
 pub struct PermissionToken {
     pub(crate) token: AuthorizationToken,
 }
 
+impl PartialEq for PermissionToken {
+    fn eq(&self, other: &Self) -> bool {
+        use AuthorizationToken::*;
+        match (&self.token, &other.token) {
+            (Primary(a), Primary(b)) => a == b,
+            (Resource(a), Resource(b)) => a == b,
+            _ => false,
+        }
+    }
+}
+
+impl Eq for PermissionToken {}
+
 impl serde::Serialize for PermissionToken {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -30,6 +43,9 @@ 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))),
+            AuthorizationToken::TokenCredential(_) => {
+                panic!("TokenCredential not supported for PermissionToken")
+            }
         };
         write!(
             f,