From 3d1647ab68e067db6b5e8548481dcc60575068fe Mon Sep 17 00:00:00 2001
From: David Mulder <dmulder@samba.org>
Date: Mon, 29 Apr 2024 11:26:24 -0600
Subject: [PATCH] Add a check for MFA in the amr of User Tokens

Signed-off-by: David Mulder <dmulder@samba.org>
---
 Cargo.toml  |  2 +-
 src/auth.rs | 65 +++++++++++++++++++++++++++++++++--------------------
 2 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 61f5cb6..9164d0d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "msal"
 description = "Microsoft Authentication Library for Rust"
-version = "0.1.18"
+version = "0.1.19"
 edition = "2021"
 authors = [
     "David Mulder <dmulder@suse.com>"
diff --git a/src/auth.rs b/src/auth.rs
index 384ed19..3bc219e 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -319,6 +319,13 @@ where
     }
 }
 
+#[derive(Deserialize, Zeroize, ZeroizeOnDrop)]
+pub struct AccessTokenPayload {
+    amr: Vec<String>,
+    tid: String,
+    upn: String,
+}
+
 #[derive(Clone, Deserialize, Zeroize, ZeroizeOnDrop)]
 pub struct UserToken {
     pub token_type: String,
@@ -355,7 +362,7 @@ impl UserToken {
         } else if let Some(access_token) = &self.access_token {
             let mut siter = access_token.splitn(3, '.');
             siter.next(); // Ignore the header
-            let payload: Value = json_from_str(
+            let payload: AccessTokenPayload = json_from_str(
                 &String::from_utf8(
                     URL_SAFE_NO_PAD
                         .decode(siter.next().ok_or_else(|| {
@@ -366,17 +373,7 @@ impl UserToken {
                 .map_err(|e| MsalError::InvalidParse(format!("{}", e)))?,
             )
             .map_err(|e| MsalError::InvalidJson(format!("{}", e)))?;
-            match payload.get("tid") {
-                Some(tid) => match tid.as_str() {
-                    Some(tid) => Ok(tid.to_string()),
-                    None => Err(MsalError::GeneralFailure(
-                        "No tid available for UserToken".to_string(),
-                    )),
-                },
-                None => Err(MsalError::GeneralFailure(
-                    "No tid available for UserToken".to_string(),
-                )),
-            }
+            Ok(payload.tid.clone())
         } else {
             Err(MsalError::GeneralFailure(
                 "No tid available for UserToken".to_string(),
@@ -408,7 +405,7 @@ impl UserToken {
                 Some(access_token) => {
                     let mut siter = access_token.splitn(3, '.');
                     siter.next(); // Ignore the header
-                    let payload: Value = json_from_str(
+                    let payload: AccessTokenPayload = json_from_str(
                         &String::from_utf8(
                             URL_SAFE_NO_PAD
                                 .decode(siter.next().ok_or_else(|| {
@@ -419,17 +416,7 @@ impl UserToken {
                         .map_err(|e| MsalError::InvalidParse(format!("{}", e)))?,
                     )
                     .map_err(|e| MsalError::InvalidJson(format!("{}", e)))?;
-                    match payload.get("upn") {
-                        Some(upn) => match upn.as_str() {
-                            Some(upn) => Ok(upn.to_string()),
-                            None => Err(MsalError::GeneralFailure(
-                                "No spn available for UserToken".to_string(),
-                            )),
-                        },
-                        None => Err(MsalError::GeneralFailure(
-                            "No spn available for UserToken".to_string(),
-                        )),
-                    }
+                    Ok(payload.upn.clone())
                 }
                 None => Err(MsalError::GeneralFailure(
                     "No spn available for UserToken".to_string(),
@@ -437,6 +424,36 @@ impl UserToken {
             },
         }
     }
+
+    /// Check if the access token amr contains an MFA authorization
+    ///
+    /// # Returns
+    ///
+    /// * Success: Whether or not the token has MFA authorization.
+    /// * Failure: An MsalError, indicating the failure.
+    pub fn amr_mfa(&self) -> Result<bool, MsalError> {
+        match &self.access_token {
+            Some(access_token) => {
+                let mut siter = access_token.splitn(3, '.');
+                siter.next(); // Ignore the header
+                let payload: AccessTokenPayload = json_from_str(
+                    &String::from_utf8(
+                        URL_SAFE_NO_PAD
+                            .decode(siter.next().ok_or_else(|| {
+                                MsalError::InvalidParse("Payload not present".to_string())
+                            })?)
+                            .map_err(|e| MsalError::InvalidBase64(format!("{}", e)))?,
+                    )
+                    .map_err(|e| MsalError::InvalidParse(format!("{}", e)))?,
+                )
+                .map_err(|e| MsalError::InvalidJson(format!("{}", e)))?;
+                Ok(payload.amr.iter().any(|s| s == "ngcmfa" || s == "mfa"))
+            }
+            None => Err(MsalError::GeneralFailure(
+                "No access token available for UserToken".to_string(),
+            )),
+        }
+    }
 }
 
 #[cfg(feature = "broker")]