From f56c5651f86ea0da43e1db02d13be87ee1363018 Mon Sep 17 00:00:00 2001
From: Xynnn007 <xynnn@linux.alibaba.com>
Date: Fri, 13 Dec 2024 16:25:37 +0800
Subject: [PATCH] ProtectedHeader: change other_fields to Map to allow nested
 fields

For JWE using EC or ECX, the ephemeral asymmetric public key should be
included in the ProtectedHeader. It will be a nested structure rather
than flatten fields.

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>
---
 src/lib.rs | 47 +++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 540a2f2..272f606 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -9,11 +9,11 @@ mod error;
 pub use error::{KbsTypesError, Result};
 
 #[cfg(all(feature = "alloc", not(feature = "std")))]
-use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};
+use alloc::{string::String, vec::Vec};
 use base64::{prelude::BASE64_URL_SAFE_NO_PAD, Engine};
-use serde_json::Value;
+use serde_json::{Map, Value};
 #[cfg(all(feature = "std", not(feature = "alloc")))]
-use std::{collections::BTreeMap, string::String};
+use std::string::String;
 
 use serde::{Deserialize, Serialize};
 
@@ -97,8 +97,8 @@ pub struct ProtectedHeader {
     pub enc: String,
 
     /// Other fields of Protected Header
-    #[serde(skip_serializing_if = "BTreeMap::is_empty", flatten)]
-    pub other_fields: BTreeMap<String, String>,
+    #[serde(skip_serializing_if = "Map::is_empty", flatten)]
+    pub other_fields: Map<String, Value>,
 }
 
 impl ProtectedHeader {
@@ -277,7 +277,7 @@ mod tests {
         let protected_header = ProtectedHeader {
             alg: "fakealg".to_string(),
             enc: "fakeenc".to_string(),
-            other_fields: BTreeMap::new(),
+            other_fields: Map::new(),
         };
 
         let aad = protected_header.generate_aad().unwrap();
@@ -311,6 +311,41 @@ mod tests {
         assert_eq!(response.aad, None);
     }
 
+    #[test]
+    fn parse_response_nested_protected_header() {
+        let data = r#"
+        {
+            "protected": "eyJhbGciOiJmYWtlYWxnIiwiZW5jIjoiZmFrZWVuYyIsImVwayI6eyJrdHkiOiJPS1AiLCJjcnYiOiJYMjU1MTkiLCJ4IjoiaFNEd0NZa3dwMVIwaTMzY3RENzNXZzJfT2cwbU9CcjA2NlNwanFxYlRtbyJ9fQo",
+            "encrypted_key": "ZmFrZWtleQ",
+            "iv": "cmFuZG9tZGF0YQ",
+            "ciphertext": "ZmFrZWVuY291dHB1dA",
+            "tag": "ZmFrZXRhZw"
+        }"#;
+
+        let response: Response = serde_json::from_str(data).unwrap();
+
+        assert_eq!(response.protected.alg, "fakealg");
+        assert_eq!(response.protected.enc, "fakeenc");
+
+        let expected_other_fields = json!({
+            "epk": {
+                "kty" : "OKP",
+                "crv": "X25519",
+                "x": "hSDwCYkwp1R0i33ctD73Wg2_Og0mOBr066SpjqqbTmo"
+            }
+        })
+        .as_object()
+        .unwrap()
+        .clone();
+
+        assert_eq!(response.protected.other_fields, expected_other_fields);
+        assert_eq!(response.encrypted_key, "fakekey".as_bytes());
+        assert_eq!(response.iv, "randomdata".as_bytes());
+        assert_eq!(response.ciphertext, "fakeencoutput".as_bytes());
+        assert_eq!(response.tag, "faketag".as_bytes());
+        assert_eq!(response.aad, None);
+    }
+
     #[test]
     fn parse_response_with_aad() {
         let data = r#"