Add serde to StackerDBMessage types

Signed-off-by: Jacinta Ferrant <[email protected]>

stacks-signer/src/client/stackerdb.rs

@@ -32,6 +32,33 @@ const SIGNATURE_SHARE_REQUEST_SLOT_ID: u32 = 7;
 const SIGNATURE_SHARE_RESPONSE_SLOT_ID: u32 = 8;
 const BLOCK_SLOT_ID: u32 = 9;
 
+/// This is required for easy serialization of the various StackerDBMessage types
+#[repr(u8)]
+enum TypePrefix {
+    Block,
+    Packet,
+}
+
+impl TypePrefix {
+    /// Convert a u8 to a TypePrefix
+    fn from_u8(value: u8) -> Option<Self> {
+        match value {
+            0 => Some(Self::Block),
+            1 => Some(Self::Packet),
+            _ => None,
+        }
+    }
+}
+
+impl From<&StackerDBMessage> for TypePrefix {
+    fn from(message: &StackerDBMessage) -> TypePrefix {
+        match message {
+            StackerDBMessage::Block(_) => TypePrefix::Block,
+            StackerDBMessage::Packet(_) => TypePrefix::Packet,
+        }
+    }
+}
+
 /// The StackerDB messages that can be sent through the .signers contract
 pub enum StackerDBMessage {
     /// The latest Nakamoto block for miners to observe
@@ -50,12 +77,39 @@ impl From<Packet> for StackerDBMessage {
 }
 
 impl StacksMessageCodec for StackerDBMessage {
-    fn consensus_serialize<W: std::io::Write>(&self, _fd: &mut W) -> Result<(), CodecError> {
-        todo!()
+    fn consensus_serialize<W: std::io::Write>(&self, fd: &mut W) -> Result<(), CodecError> {
+        fd.write_all(&[TypePrefix::from(self) as u8])
+            .map_err(CodecError::WriteError)?;
+        match self {
+            StackerDBMessage::Packet(packet) => {
+                let message_bytes = bincode::serialize(&packet)
+                    .map_err(|e| CodecError::SerializeError(e.to_string()))?;
+                message_bytes.consensus_serialize(fd)
+            }
+            StackerDBMessage::Block(block) => block.consensus_serialize(fd),
+        }
     }
 
-    fn consensus_deserialize<R: std::io::Read>(_fd: &mut R) -> Result<Self, CodecError> {
-        todo!()
+    fn consensus_deserialize<R: std::io::Read>(fd: &mut R) -> Result<Self, CodecError> {
+        let mut prefix = [0];
+        fd.read_exact(&mut prefix)
+            .map_err(|e| CodecError::DeserializeError(e.to_string()))?;
+        let prefix = TypePrefix::from_u8(prefix[0]).ok_or(CodecError::DeserializeError(
+            "Bad StackerDBMessage prefix".into(),
+        ))?;
+
+        match prefix {
+            TypePrefix::Packet => {
+                let message_bytes = Vec::<u8>::consensus_deserialize(fd)?;
+                let packet = bincode::deserialize(&message_bytes)
+                    .map_err(|e| CodecError::DeserializeError(e.to_string()))?;
+                Ok(Self::Packet(packet))
+            }
+            TypePrefix::Block => {
+                let block = NakamotoBlock::consensus_deserialize(fd)?;
+                Ok(StackerDBMessage::Block(block))
+            }
+        }
     }
 }
 
@@ -162,4 +216,96 @@ impl StackerDB {
 }
 
 #[cfg(test)]
-mod tests {}
+mod tests {
+    use blockstack_lib::chainstate::nakamoto::{NakamotoBlock, NakamotoBlockHeader};
+    use blockstack_lib::chainstate::stacks::StacksTransaction;
+    use rand_core::OsRng;
+    use stacks_common::codec::StacksMessageCodec;
+    use stacks_common::types::chainstate::{ConsensusHash, StacksBlockId, TrieHash};
+    use stacks_common::util::hash::{MerkleTree, Sha512Trunc256Sum};
+    use stacks_common::util::secp256k1::{MessageSignature, SchnorrSignature};
+    use wsts::curve::scalar::Scalar;
+    use wsts::net::{Message, Packet, Signable, SignatureShareRequest};
+
+    use super::StackerDBMessage;
+
+    #[test]
+    fn serde_stackerdb_message_block() {
+        let txs: Vec<StacksTransaction> = vec![];
+        let mut header = NakamotoBlockHeader {
+            version: 1,
+            chain_length: 2,
+            burn_spent: 3,
+            consensus_hash: ConsensusHash([0x04; 20]),
+            parent_block_id: StacksBlockId([0x05; 32]),
+            tx_merkle_root: Sha512Trunc256Sum([0x06; 32]),
+            state_index_root: TrieHash([0x07; 32]),
+            miner_signature: MessageSignature::empty(),
+            signer_signature: SchnorrSignature::default(),
+        };
+        let txid_vecs = txs.iter().map(|tx| tx.txid().as_bytes().to_vec()).collect();
+
+        let merkle_tree = MerkleTree::<Sha512Trunc256Sum>::new(&txid_vecs);
+        let tx_merkle_root = merkle_tree.root();
+
+        header.tx_merkle_root = tx_merkle_root;
+
+        let block = NakamotoBlock { header, txs };
+
+        let msg = StackerDBMessage::Block(block.clone());
+        let serialized_bytes = msg.serialize_to_vec();
+        let deserialized_msg =
+            StackerDBMessage::consensus_deserialize(&mut &serialized_bytes[..]).unwrap();
+        match deserialized_msg {
+            StackerDBMessage::Block(deserialized_block) => {
+                assert_eq!(deserialized_block, block);
+            }
+            _ => panic!("Wrong message type. Expected StackerDBMessage::Block"),
+        }
+    }
+
+    #[test]
+    fn serde_stackerdb_message_packet() {
+        let mut rng = OsRng;
+        let private_key = Scalar::random(&mut rng);
+        let to_sign = "One, two, three, four, five? That's amazing. I've got the same combination on my luggage.".as_bytes();
+        let sig_share_request = SignatureShareRequest {
+            dkg_id: 1,
+            sign_id: 5,
+            sign_iter_id: 4,
+            nonce_responses: vec![],
+            message: to_sign.to_vec(),
+            is_taproot: false,
+            merkle_root: None,
+        };
+        let packet = Packet {
+            sig: sig_share_request
+                .sign(&private_key)
+                .expect("Failed to sign SignatureShareRequest"),
+            msg: Message::SignatureShareRequest(sig_share_request),
+        };
+
+        let msg = StackerDBMessage::Packet(packet.clone());
+        let serialized_bytes = msg.serialize_to_vec();
+        let deserialized_msg =
+            StackerDBMessage::consensus_deserialize(&mut &serialized_bytes[..]).unwrap();
+        match deserialized_msg {
+            StackerDBMessage::Packet(deserialized_packet) => {
+                assert_eq!(deserialized_packet.sig, packet.sig);
+                match deserialized_packet.msg {
+                    Message::SignatureShareRequest(deserialized_message) => {
+                        assert_eq!(deserialized_message.dkg_id, 1);
+                        assert_eq!(deserialized_message.sign_id, 5);
+                        assert_eq!(deserialized_message.sign_iter_id, 4);
+                        assert!(deserialized_message.nonce_responses.is_empty());
+                        assert_eq!(deserialized_message.message.as_slice(), to_sign);
+                        assert!(!deserialized_message.is_taproot);
+                        assert!(deserialized_message.merkle_root.is_none());
+                    }
+                    _ => panic!("Wrong message type. Expected Message::SignatureShareRequest"),
+                }
+            }
+            _ => panic!("Wrong message type. Expected StackerDBMessage::Packet."),
+        }
+    }
+}