From f313edbb143bda77ed9d9f05d4191aa25916bdd1 Mon Sep 17 00:00:00 2001
From: Xynnn007 <xynnn@linux.alibaba.com>
Date: Wed, 27 Nov 2024 10:02:28 +0800
Subject: [PATCH] Define ProtectedHeader for Response

The struct of the ProtectedHeader is actually well defined in RFC7519,
the changing from string to a struct would do help to serialization and
deserialization.

See https://datatracker.ietf.org/doc/html/rfc7516

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

diff --git a/src/lib.rs b/src/lib.rs
index 4020f71..647b80e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,7 @@ extern crate alloc;
 #[cfg(all(feature = "alloc", not(feature = "std")))]
 use alloc::string::String;
 use serde_json::Value;
+use std::collections::BTreeMap;
 #[cfg(feature = "std")]
 use std::string::String;
 
@@ -85,9 +86,21 @@ pub struct Attestation {
     pub tee_evidence: Value,
 }
 
+#[derive(Clone, Serialize, Deserialize, Debug)]
+pub struct ProtectedHeader {
+    /// Enryption algorithm for encrypted key
+    alg: String,
+    /// Encryption algorithm for payload
+    enc: String,
+
+    /// Other fields of Protected Header
+    #[serde(skip_serializing_if = "BTreeMap::is_empty", flatten)]
+    other_fields: BTreeMap<String, String>,
+}
+
 #[derive(Clone, Serialize, Deserialize, Debug)]
 pub struct Response {
-    pub protected: String,
+    pub protected: ProtectedHeader,
     pub encrypted_key: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub aad: Option<String>,
@@ -141,7 +154,10 @@ mod tests {
     fn parse_response() {
         let data = r#"
         {
-            "protected": "fakejoseheader",
+            "protected": {
+                "alg": "fakealg",
+                "enc": "fakeenc"
+            },
             "encrypted_key": "fakekey",
             "iv": "randomdata",
             "ciphertext": "fakeencoutput",
@@ -150,7 +166,9 @@ mod tests {
 
         let response: Response = serde_json::from_str(data).unwrap();
 
-        assert_eq!(response.protected, "fakejoseheader");
+        assert_eq!(response.protected.alg, "fakealg");
+        assert_eq!(response.protected.enc, "fakeenc");
+        assert!(response.protected.other_fields.is_empty());
         assert_eq!(response.encrypted_key, "fakekey");
         assert_eq!(response.iv, "randomdata");
         assert_eq!(response.ciphertext, "fakeencoutput");
@@ -162,7 +180,38 @@ mod tests {
     fn parse_response_with_aad() {
         let data = r#"
         {
-            "protected": "fakejoseheader",
+            "protected": {
+                "alg": "fakealg",
+                "enc": "fakeenc"
+            },
+            "encrypted_key": "fakekey",
+            "iv": "randomdata",
+            "aad": "fakeaad",
+            "ciphertext": "fakeencoutput",
+            "tag": "faketag"
+        }"#;
+
+        let response: Response = serde_json::from_str(data).unwrap();
+
+        assert_eq!(response.protected.alg, "fakealg");
+        assert_eq!(response.protected.enc, "fakeenc");
+        assert!(response.protected.other_fields.is_empty());
+        assert_eq!(response.encrypted_key, "fakekey");
+        assert_eq!(response.iv, "randomdata");
+        assert_eq!(response.ciphertext, "fakeencoutput");
+        assert_eq!(response.tag, "faketag");
+        assert_eq!(response.aad, Some("fakeaad".into()));
+    }
+
+    #[test]
+    fn parse_response_with_protectedheader() {
+        let data = r#"
+        {
+            "protected": {
+                "alg": "fakealg",
+                "enc": "fakeenc",
+                "fakefield": "fakevalue"
+            },
             "encrypted_key": "fakekey",
             "iv": "randomdata",
             "aad": "fakeaad",
@@ -172,7 +221,9 @@ mod tests {
 
         let response: Response = serde_json::from_str(data).unwrap();
 
-        assert_eq!(response.protected, "fakejoseheader");
+        assert_eq!(response.protected.alg, "fakealg");
+        assert_eq!(response.protected.enc, "fakeenc");
+        assert_eq!(response.protected.other_fields["fakefield"], "fakevalue");
         assert_eq!(response.encrypted_key, "fakekey");
         assert_eq!(response.iv, "randomdata");
         assert_eq!(response.ciphertext, "fakeencoutput");