From d505a80b86e27a78ff27a5385a4b2048a8852937 Mon Sep 17 00:00:00 2001
From: David Mulder <dmulder@samba.org>
Date: Thu, 1 Feb 2024 09:37:31 -0700
Subject: [PATCH] Derive the decryption key using the ctx

As explained in [MS-OAPXBC] 3.1.5.1.3.3. The
decryption algorithm is incorrect in the Jwe
header when the response is a PrtV2, and instead
should be decrypted with aes_256_cbc. This is not
documented by Microsoft, but has just been
observed in practice.

Signed-off-by: David Mulder <dmulder@samba.org>
---
 Cargo.toml              |  1 +
 src/compact.rs          |  2 ++
 src/crypto/a256gcm.rs   |  2 +-
 src/crypto/ms_oapxbc.rs | 56 +++++++++++++++++++++++++++++++++++++----
 4 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2730ee4..e9f75eb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,6 +32,7 @@ url = { version = "^2.2.2", features = ["serde"] }
 uuid = { version = "^1.0.0", features = ["serde"] }
 tracing = "^0.1.34"
 hex = "0.4"
+openssl-kdf = "0.4.2"
 
 
 [dev-dependencies]
diff --git a/src/compact.rs b/src/compact.rs
index 90abd38..e049bf8 100644
--- a/src/compact.rs
+++ b/src/compact.rs
@@ -365,6 +365,8 @@ pub struct JweProtectedHeader {
         skip_serializing_if = "Option::is_none"
     )]
     pub(crate) x5t_s256: Option<()>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) ctx: Option<String>,
     // Don't allow extra header names?
 }
 
diff --git a/src/crypto/a256gcm.rs b/src/crypto/a256gcm.rs
index 233c6fe..93aa635 100644
--- a/src/crypto/a256gcm.rs
+++ b/src/crypto/a256gcm.rs
@@ -16,7 +16,7 @@ const AUTH_TAG_LEN: usize = 16;
 
 #[derive(Clone)]
 pub struct JweA256GCMEncipher {
-    aes_key: [u8; KEY_LEN],
+    pub(crate) aes_key: [u8; KEY_LEN],
 }
 
 #[cfg(test)]
diff --git a/src/crypto/ms_oapxbc.rs b/src/crypto/ms_oapxbc.rs
index 5fa0455..4247d0d 100644
--- a/src/crypto/ms_oapxbc.rs
+++ b/src/crypto/ms_oapxbc.rs
@@ -11,9 +11,15 @@ use super::rsaes_oaep::{JweRSAOAEPDecipher, JweRSAOAEPEncipher};
 
 use crate::jwe::JweBuilder;
 
+use openssl::hash::MessageDigest;
 use openssl::pkey::Private;
 use openssl::pkey::Public;
 use openssl::rsa::Rsa;
+use openssl::symm::decrypt;
+use openssl::symm::Cipher;
+use openssl_kdf::{perform_kdf, KdfArgument, KdfKbMode, KdfMacType, KdfType};
+
+use base64::{engine::general_purpose, Engine as _};
 
 /// A [MS-OAPXBC] 3.2.5.1.2.2 yielded session key. This is used as a form of key agreement
 /// for MS clients, where this key can now be used to encipher and decipher arbitrary
@@ -28,6 +34,11 @@ pub enum MsOapxbcSessionKey {
     },
 }
 
+pub enum PrtVersion {
+    V2,
+    V3,
+}
+
 #[cfg(test)]
 impl MsOapxbcSessionKey {
     pub(crate) fn assert_key(&self, key: &[u8]) -> bool {
@@ -91,7 +102,7 @@ impl MsOapxbcSessionKey {
 
 impl MsOapxbcSessionKey {
     /// Given a JWE in compact form, decipher and authenticate its content.
-    pub fn decipher(&self, jwec: &JweCompact) -> Result<Jwe, JwtError> {
+    pub fn decipher(&self, jwec: &JweCompact, prt_vers: PrtVersion) -> Result<Jwe, JwtError> {
         // Alg must be direct.
         if jwec.header.alg != JweAlg::DIRECT {
             return Err(JwtError::AlgorithmUnavailable);
@@ -104,10 +115,45 @@ impl MsOapxbcSessionKey {
                     return Err(JwtError::CipherUnavailable);
                 }
 
-                aes_key.decipher_inner(jwec).map(|payload| Jwe {
-                    header: jwec.header.clone(),
-                    payload,
-                })
+                // If a ctx is present in the header, derive the session key
+                let mut session_key = aes_key.clone();
+                if let Some(ctx) = &jwec.header.ctx {
+                    let decoded_ctx = general_purpose::STANDARD
+                        .decode(ctx)
+                        .map_err(|_| JwtError::InvalidBase64)?;
+                    let args = [
+                        &KdfArgument::KbMode(KdfKbMode::Counter),
+                        &KdfArgument::Mac(KdfMacType::Hmac(MessageDigest::sha256())),
+                        &KdfArgument::Salt(b"AzureAD-SecureConversation"),
+                        &KdfArgument::KbInfo(&decoded_ctx),
+                        &KdfArgument::Key(&session_key.aes_key),
+                    ];
+                    let derived_key =
+                        perform_kdf(KdfType::KeyBased, &args, session_key.aes_key.len())
+                            .map_err(|_| JwtError::InvalidKey)?;
+                    session_key = JweA256GCMEncipher::try_from(derived_key.as_slice())?;
+                }
+
+                match prt_vers {
+                    PrtVersion::V2 => {
+                        let cipher = Cipher::aes_256_cbc();
+                        let decrypted = decrypt(
+                            cipher,
+                            &session_key.aes_key,
+                            Some(&jwec.iv),
+                            &jwec.ciphertext,
+                        )
+                        .map_err(|_| JwtError::OpenSSLError)?;
+                        Ok(Jwe {
+                            header: jwec.header.clone(),
+                            payload: decrypted,
+                        })
+                    }
+                    PrtVersion::V3 => aes_key.decipher_inner(jwec).map(|payload| Jwe {
+                        header: jwec.header.clone(),
+                        payload,
+                    }),
+                }
             }
         }
     }