Add proper version resolution for MessageKit, RetrievalKit, and Reenc…

…ryptionRequest
nucypher · Sep 14, 2022 · 21f657f · 21f657f
1 parent 82cbb7b
commit 21f657f
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 15 deletions.
diff --git a/nucypher-core/src/message_kit.rs b/nucypher-core/src/message_kit.rs
@@ -67,24 +67,38 @@ impl MessageKit {
     }
 }
 
+#[derive(Deserialize)]
+struct MessageKit0 {
+    capsule: Capsule,
+    #[serde(with = "serde_bytes::as_base64")]
+    ciphertext: Box<[u8]>,
+}
+
 impl<'a> ProtocolObjectInner<'a> for MessageKit {
     fn brand() -> [u8; 4] {
         *b"MKit"
     }
 
     fn version() -> (u16, u16) {
-        (1, 0)
+        (1, 1)
     }
 
     fn unversioned_to_bytes(&self) -> Box<[u8]> {
         messagepack_serialize(&self)
     }
 
     fn unversioned_from_bytes(minor_version: u16, bytes: &[u8]) -> Option<Result<Self, String>> {
-        if minor_version == 0 {
-            Some(messagepack_deserialize(bytes))
-        } else {
-            None
+        match minor_version {
+            0 => {
+                let mkit = messagepack_deserialize::<MessageKit0>(bytes).map(|mkit| MessageKit {
+                    capsule: mkit.capsule,
+                    ciphertext: mkit.ciphertext,
+                    conditions: None,
+                });
+                Some(mkit)
+            }
+            1 => Some(messagepack_deserialize(bytes)),
+            _ => None,
         }
     }
 }

diff --git a/nucypher-core/src/reencryption.rs b/nucypher-core/src/reencryption.rs
@@ -57,24 +57,46 @@ impl ReencryptionRequest {
     }
 }
 
+#[derive(Deserialize)]
+struct ReencryptionRequest0 {
+    capsules: Box<[Capsule]>,
+    hrac: HRAC,
+    encrypted_kfrag: EncryptedKeyFrag,
+    publisher_verifying_key: PublicKey,
+    bob_verifying_key: PublicKey,
+}
+
 impl<'a> ProtocolObjectInner<'a> for ReencryptionRequest {
     fn brand() -> [u8; 4] {
         *b"ReRq"
     }
 
     fn version() -> (u16, u16) {
-        (1, 0)
+        (1, 1)
     }
 
     fn unversioned_to_bytes(&self) -> Box<[u8]> {
         messagepack_serialize(&self)
     }
 
     fn unversioned_from_bytes(minor_version: u16, bytes: &[u8]) -> Option<Result<Self, String>> {
-        if minor_version == 0 {
-            Some(messagepack_deserialize(bytes))
-        } else {
-            None
+        match minor_version {
+            0 => {
+                let req = messagepack_deserialize::<ReencryptionRequest0>(bytes).map(|req| {
+                    ReencryptionRequest {
+                        capsules: req.capsules,
+                        hrac: req.hrac,
+                        encrypted_kfrag: req.encrypted_kfrag,
+                        publisher_verifying_key: req.publisher_verifying_key,
+                        bob_verifying_key: req.bob_verifying_key,
+                        conditions: None,
+                        context: None,
+                    }
+                });
+                Some(req)
+            }
+            1 => Some(messagepack_deserialize(bytes)),
+            _ => None,
         }
     }
 }

diff --git a/nucypher-core/src/retrieval_kit.rs b/nucypher-core/src/retrieval_kit.rs
@@ -50,24 +50,38 @@ impl RetrievalKit {
     }
 }
 
+#[derive(Deserialize)]
+struct RetrievalKit0 {
+    capsule: Capsule,
+    queried_addresses: BTreeSet<Address>,
+}
+
 impl<'a> ProtocolObjectInner<'a> for RetrievalKit {
     fn brand() -> [u8; 4] {
         *b"RKit"
     }
 
     fn version() -> (u16, u16) {
-        (1, 0)
+        (1, 1)
     }
 
     fn unversioned_to_bytes(&self) -> Box<[u8]> {
         messagepack_serialize(&self)
     }
 
     fn unversioned_from_bytes(minor_version: u16, bytes: &[u8]) -> Option<Result<Self, String>> {
-        if minor_version == 0 {
-            Some(messagepack_deserialize(bytes))
-        } else {
-            None
+        match minor_version {
+            0 => {
+                let rkit =
+                    messagepack_deserialize::<RetrievalKit0>(bytes).map(|rkit| RetrievalKit {
+                        capsule: rkit.capsule,
+                        queried_addresses: rkit.queried_addresses,
+                        conditions: None,
+                    });
+                Some(rkit)
+            }
+            1 => Some(messagepack_deserialize(bytes)),
+            _ => None,
         }
     }
 }