From 1eb7aeab42922d077f87272431f279696ec1f492 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Thu, 14 Nov 2024 00:04:36 +0800
Subject: [PATCH] feat: update jose-test

---
 __crypto/jose-test/src/jose.rs | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/__crypto/jose-test/src/jose.rs b/__crypto/jose-test/src/jose.rs
index 5760c6e..87b0b15 100644
--- a/__crypto/jose-test/src/jose.rs
+++ b/__crypto/jose-test/src/jose.rs
@@ -92,7 +92,10 @@ pub fn serialize_jwe_rsa(payload: &[u8], jwk: &Jwk) -> XResult<String> {
 }
 
 pub fn deserialize_jwe_rsa_2(jwe: &str, rsa: &RsaPrivateKey) -> XResult<(Vec<u8>, JweHeader2)> {
-    deserialize_jwe_fn(jwe, |key_wrap| -> XResult<(Vec<u8>)> {
+    deserialize_jwe_fn(jwe, |alg, key_wrap| -> XResult<Vec<u8>> {
+        if alg != "RSA-OAEP" {
+            return simple_error!("Invalid JWE header alg: {}", alg);
+        }
         Ok(opt_result!(rsa.decrypt(Oaep::new::<Sha1>(), &key_wrap), "Unwrap key failed: {}"))
     })
 }
@@ -131,7 +134,10 @@ pub fn serialize_jwe_aes(payload: &[u8], key: &[u8]) -> XResult<String> {
 }
 
 pub fn deserialize_jwe_aes_2(jwe: &str, key: [u8; 32]) -> XResult<(Vec<u8>, JweHeader2)> {
-    deserialize_jwe_fn(jwe, |key_wrap| -> XResult<(Vec<u8>)> {
+    deserialize_jwe_fn(jwe, |alg, key_wrap| -> XResult<Vec<u8>> {
+        if alg != "A256KW" {
+            return simple_error!("Invalid JWE header alg: {}", alg);
+        }
         let kek = Kek::from(key);
         Ok(opt_result!(kek.unwrap_vec(&key_wrap), "Unwrap key failed: {}"))
     })
@@ -175,7 +181,7 @@ where
 
 fn deserialize_jwe_fn<F>(jwe: &str, key_unwrap_fn: F) -> XResult<(Vec<u8>, JweHeader2)>
 where
-    F: Fn(&[u8]) -> XResult<Vec<u8>>,
+    F: Fn(&str, &[u8]) -> XResult<Vec<u8>>,
 {
     let jwe_parts = jwe.split(".").collect::<Vec<&str>>();
     if jwe_parts.len() != 5 {
@@ -183,12 +189,16 @@ where
     }
     let header_bytes = opt_result!(decode_url_safe_no_pad(jwe_parts[0]), "Invalid JWE header: {}, JWE: {}", jwe);
     let header: JweHeader2 = opt_result!(serde_json::from_slice(&header_bytes), "Invalid JWE header: {}, JWE: {}", jwe);
+    if header.enc != "A256GCM" {
+        return simple_error!("Invalid JWE header enc: {}", header.enc);
+    }
+
     let cek = opt_result!(decode_url_safe_no_pad(jwe_parts[1]), "Invalid JWE CEK: {}, JWE: {}", jwe);
     let iv = opt_result!(decode_url_safe_no_pad(jwe_parts[2]), "Invalid JWE IV: {}, JWE: {}", jwe);
     let ciphertext = opt_result!(decode_url_safe_no_pad(jwe_parts[3]), "Invalid JWE ciphertext: {}, JWE: {}", jwe);
     let tag = opt_result!(decode_url_safe_no_pad(jwe_parts[4]), "Invalid JWE tag: {}, JWE: {}", jwe);
 
-    let data_key = key_unwrap_fn(&cek)?;
+    let data_key = key_unwrap_fn(&header.alg, &cek)?;
     let data_key_b32 = opt_result!(to_bytes32(&data_key), "Invalid JWE CEK: {}, JWE: {}", jwe);
 
     let mut decryptor = Aes256GcmStreamDecryptor::new(data_key_b32, &iv);