Fixup

common/src/client_commands.rs

@@ -1,13 +1,26 @@
 // Vanadium VM client commands (responsed to InterruptedExecution status word), and other related types
 
+use crate::constants::PAGE_SIZE;
 use alloc::vec::Vec;
 
 #[cfg(feature = "device_sdk")]
 use ledger_device_sdk::io::Comm;
 
 pub trait Message: Sized {
+    fn serialize_with<F: FnMut(&[u8])>(&self, f: F);
+
     #[cfg(feature = "device_sdk")]
-    fn serialize_to_comm(&self, comm: &mut Comm);
+    #[inline]
+    fn serialize_to_comm(&self, comm: &mut Comm) {
+        self.serialize_with(|data| comm.append(data));
+    }
+
+    fn serialize(&self) -> Vec<u8> {
+        let mut result = Vec::new();
+        self.serialize_with(|data| result.extend_from_slice(data));
+        result
+    }
+
     fn deserialize(data: &[u8]) -> Result<Self, &'static str>;
 }
 
@@ -57,6 +70,130 @@ impl TryFrom<u8> for SectionKind {
     }
 }
 
+// We use the _Message ending for messages from the VM to the host, and the _Response ending for messages from the host to the VM.
+
+pub struct CommitPageMessage {
+    pub command_code: ClientCommandCode,
+    pub section_kind: SectionKind,
+    pub page_index: u32,
+}
+
+impl CommitPageMessage {
+    #[inline]
+    pub fn new(section_kind: SectionKind, page_index: u32) -> Self {
+        CommitPageMessage {
+            command_code: ClientCommandCode::CommitPage,
+            section_kind,
+            page_index,
+        }
+    }
+}
+
+impl Message for CommitPageMessage {
+    #[inline]
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&[self.command_code as u8]);
+        f(&[self.section_kind as u8]);
+        f(&self.page_index.to_be_bytes());
+    }
+
+    fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
+        let command_code = ClientCommandCode::try_from(data[0])?;
+        if !matches!(command_code, ClientCommandCode::CommitPage) || data.len() != 6 {
+            return Err("Invalid data for CommitPageMessage");
+        }
+        let section_kind = SectionKind::try_from(data[1])?;
+        let page_index = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+
+        Ok(CommitPageMessage {
+            command_code,
+            section_kind,
+            page_index,
+        })
+    }
+}
+
+pub struct CommitPageContentMessage {
+    pub command_code: ClientCommandCode,
+    pub data: Vec<u8>,
+}
+
+impl CommitPageContentMessage {
+    #[inline]
+    pub fn new(data: Vec<u8>) -> Self {
+        if data.len() != PAGE_SIZE {
+            panic!("Invalid data length for CommitPageContentMessage");
+        }
+        CommitPageContentMessage {
+            command_code: ClientCommandCode::CommitPageContent,
+            data,
+        }
+    }
+}
+
+impl Message for CommitPageContentMessage {
+    #[inline]
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&[self.command_code as u8]);
+        f(&self.data);
+    }
+
+    fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
+        let command_code = ClientCommandCode::try_from(data[0])?;
+        if !matches!(command_code, ClientCommandCode::CommitPageContent) {
+            return Err("Invalid data for CommitPageContentMessage");
+        }
+        if data.len() != PAGE_SIZE + 1 {
+            return Err("Invalid data for CommitPageContentMessage");
+        }
+        Ok(CommitPageContentMessage {
+            command_code,
+            data: data[1..].to_vec(),
+        })
+    }
+}
+
+pub struct GetPageMessage {
+    pub command_code: ClientCommandCode,
+    pub section_kind: SectionKind,
+    pub page_index: u32,
+}
+
+impl GetPageMessage {
+    #[inline]
+    pub fn new(section_kind: SectionKind, page_index: u32) -> Self {
+        GetPageMessage {
+            command_code: ClientCommandCode::GetPage,
+            section_kind,
+            page_index,
+        }
+    }
+}
+
+impl Message for GetPageMessage {
+    #[inline]
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&[self.command_code as u8]);
+        f(&[self.section_kind as u8]);
+        f(&self.page_index.to_be_bytes());
+    }
+
+    fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
+        let command_code = ClientCommandCode::try_from(data[0])?;
+        if !matches!(command_code, ClientCommandCode::GetPage) || data.len() != 6 {
+            return Err("Invalid data for GetPageMessage");
+        }
+        let section_kind = SectionKind::try_from(data[1])?;
+        let page_index = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+
+        Ok(GetPageMessage {
+            command_code,
+            section_kind,
+            page_index,
+        })
+    }
+}
+
 pub struct SendBufferMessage {
     pub command_code: ClientCommandCode,
     pub total_remaining_size: u32,
@@ -79,12 +216,11 @@ impl SendBufferMessage {
 }
 
 impl Message for SendBufferMessage {
-    #[cfg(feature = "device_sdk")]
     #[inline]
-    fn serialize_to_comm(&self, comm: &mut Comm) {
-        comm.append(&[self.command_code as u8]);
-        comm.append(&self.total_remaining_size.to_be_bytes());
-        comm.append(&self.data);
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&[self.command_code as u8]);
+        f(&self.total_remaining_size.to_be_bytes());
+        f(&self.data);
     }
 
     fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
@@ -121,12 +257,10 @@ impl ReceiveBufferMessage {
 }
 
 impl Message for ReceiveBufferMessage {
-    #[cfg(feature = "device_sdk")]
     #[inline]
-    fn serialize_to_comm(&self, comm: &mut Comm) {
-        comm.append(&[self.command_code as u8]);
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&[self.command_code as u8]);
     }
-
     fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
         let command_code = ClientCommandCode::try_from(data[0])?;
         if !matches!(command_code, ClientCommandCode::ReceiveBuffer) || data.len() != 1 {
@@ -136,3 +270,41 @@ impl Message for ReceiveBufferMessage {
         Ok(ReceiveBufferMessage { command_code })
     }
 }
+
+pub struct ReceiveBufferResponse {
+    pub remaining_length: u32,
+    pub content: Vec<u8>,
+}
+
+impl ReceiveBufferResponse {
+    #[inline]
+    pub fn new(remaining_length: u32, content: Vec<u8>) -> Self {
+        ReceiveBufferResponse {
+            remaining_length,
+            content,
+        }
+    }
+}
+
+impl Message for ReceiveBufferResponse {
+    #[inline]
+    fn serialize_with<F: FnMut(&[u8])>(&self, mut f: F) {
+        f(&self.remaining_length.to_be_bytes());
+        f(&self.content);
+    }
+
+    #[inline]
+    fn deserialize(data: &[u8]) -> Result<Self, &'static str> {
+        if data.len() < 4 {
+            return Err("Invalid data for ReceiveBufferResponse");
+        }
+        let remaining_length = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
+        if data.len() - 4 > remaining_length as usize {
+            return Err("Data for ReceiveBufferResponse is too long");
+        }
+        Ok(ReceiveBufferResponse {
+            remaining_length,
+            content: data[4..].to_vec(),
+        })
+    }
+}

host/src/vanadium_client.rs

@@ -3,7 +3,8 @@ use std::io::{stdin, stdout, Write};
 
 use common::accumulator::{HashOutput, Hasher, MerkleAccumulator, VectorAccumulator};
 use common::client_commands::{
-    ClientCommandCode, Message, ReceiveBufferMessage, SectionKind, SendBufferMessage,
+    ClientCommandCode, CommitPageContentMessage, CommitPageMessage, GetPageMessage, Message,
+    ReceiveBufferMessage, ReceiveBufferResponse, SectionKind, SendBufferMessage,
 };
 use common::constants::{page_start, PAGE_SIZE};
 use common::manifest::Manifest;
@@ -157,9 +158,11 @@ impl<'a> VAppEngine<'a> {
         transport: &T,
         command: &[u8],
     ) -> Result<(StatusWord, Vec<u8>), &'static str> {
-        let section_kind: SectionKind = command[1].try_into()?;
-
-        let page_index = u32::from_be_bytes([command[2], command[3], command[4], command[5]]);
+        let GetPageMessage {
+            command_code: _,
+            section_kind,
+            page_index,
+        } = GetPageMessage::deserialize(command)?;
 
         let segment = match section_kind {
             SectionKind::Code => &self.code_seg,
@@ -190,19 +193,16 @@ impl<'a> VAppEngine<'a> {
         transport: &T,
         command: &[u8],
     ) -> Result<(StatusWord, Vec<u8>), &'static str> {
-        let section_kind: SectionKind = command[1].try_into()?;
-        let page_index = u32::from_be_bytes([command[2], command[3], command[4], command[5]]);
+        let msg = CommitPageMessage::deserialize(command)?;
 
-        let segment = match section_kind {
+        let segment = match msg.section_kind {
             SectionKind::Code => {
                 return Err("The code segment is immutable");
             }
             SectionKind::Data => &mut self.data_seg,
             SectionKind::Stack => &mut self.stack_seg,
         };
 
-        println!("Committing page {}. Requesting content.", page_index);
-
         // get the next message, which contains the content of the page
         let (tmp_status, tmp_result) = transport
             .exchange(&apdu_continue(vec![]))
@@ -212,18 +212,15 @@ impl<'a> VAppEngine<'a> {
         if tmp_status != StatusWord::InterruptedExecution {
             return Err("Expected InterruptedExecution status word");
         }
-        if tmp_result[0] != ClientCommandCode::CommitPageContent as u8 {
-            return Err("Expected CommitPageContent client command");
-        }
-        if tmp_result.len() != 1 + PAGE_SIZE {
-            return Err("Invalid page content length");
-        }
 
-        let update_proof = segment.store_page(page_index, &tmp_result[1..])?;
+        let CommitPageContentMessage {
+            command_code: _,
+            data,
+        } = CommitPageContentMessage::deserialize(&tmp_result)?;
 
-        // TODO: for now we ignore the update proof
+        let update_proof = segment.store_page(msg.page_index, &data)?;
 
-        println!("Updated; should send the proof, but we send an empty message instead");
+        // TODO: for now we ignore the update proof
 
         Ok(transport
             .exchange(&apdu_continue(vec![]))
@@ -243,13 +240,6 @@ impl<'a> VAppEngine<'a> {
             data: mut buf,
         } = SendBufferMessage::deserialize(command)?;
 
-        println!(
-            "Receiving buffer of length: {}. Length of first chunk: {}. First chunk: {:?}",
-            remaining_len,
-            buf.len(),
-            buf
-        );
-
         if (buf.len() as u32) > remaining_len {
             return Err("Received data length exceeds expected remaining length");
         }
@@ -313,9 +303,11 @@ impl<'a> VAppEngine<'a> {
         loop {
             // TODO: wrong if the buffer is long
             let chunk_len = min(remaining_len, 255 - 4);
-            let mut data = vec![];
-            data.extend_from_slice(&remaining_len.to_be_bytes());
-            data.extend_from_slice(&bytes[offset..offset + chunk_len as usize]);
+            let data = ReceiveBufferResponse::new(
+                remaining_len,
+                bytes[offset..offset + chunk_len as usize].to_vec(),
+            )
+            .serialize();
 
             let (status, result) = self
                 .exchange_and_process_page_requests(transport, &apdu_continue(data))
@@ -353,12 +345,10 @@ impl<'a> VAppEngine<'a> {
 
         loop {
             if status != StatusWord::InterruptedExecution {
-                println!("Unexpected status word: {:?}", status);
                 return Err("Unexpected status word");
             }
 
             if result.len() == 0 {
-                println!("Unexpected empty command");
                 return Err("empty command");
             }