[DRAFT] Implement hmac for repeated pages

Incomplete: using a placeholder instead of the actual HMAC
(that should use a private key, possibly key deterministically
crated from the app hash (this would allow stored HMACs to
remain valid for future executions of the app).

Exploratory commit to modify the protocol to compute an hmac
before dropping a code page, so that the required data for pages
sent multiple times is reduced; moreover, the client could (in
principle) store it permanently for future executions of the
V-App).

Author: bigspider

client-sdk/src/vanadium_client.rs

@@ -18,8 +18,8 @@ use common::client_commands::{
     BufferType, ClientCommandCode, CommitPageMessage, CommitPageProofContinuedMessage,
     CommitPageProofContinuedResponse, CommitPageProofResponse, GetPageMessage,
     GetPageProofContinuedMessage, GetPageProofContinuedResponse, GetPageResponse, Message,
-    MessageDeserializationError, ReceiveBufferMessage, ReceiveBufferResponse, SectionKind,
-    SendBufferContinuedMessage, SendBufferMessage,
+    MessageDeserializationError, PageProofKind, ReceiveBufferMessage, ReceiveBufferResponse,
+    SectionKind, SendBufferContinuedMessage, SendBufferMessage, StorePageProofMessage,
 };
 use common::constants::{DEFAULT_STACK_START, PAGE_SIZE};
 use common::manifest::Manifest;
@@ -128,6 +128,7 @@ struct VAppEngine<E: std::fmt::Debug + Send + Sync + 'static> {
     engine_to_client_sender: mpsc::Sender<VAppMessage>,
     client_to_engine_receiver: mpsc::Receiver<ClientMessage>,
     print_writer: Box<dyn std::io::Write + Send>,
+    stored_code_page_hmacs: Vec<Option<[u8; 32]>>,
 }
 
 impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
@@ -166,6 +167,7 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
             (status, result) = match client_command_code {
                 ClientCommandCode::GetPage => self.process_get_page(&result).await?,
                 ClientCommandCode::CommitPage => self.process_commit_page(&result).await?,
+                ClientCommandCode::StorePageProof => self.process_store_page_proof(&result).await?,
                 _ => return Ok((status, result)),
             }
         }
@@ -204,10 +206,37 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
         let is_encrypted = header[0] != 0;
         let nonce: [u8; 12] = header[1..13].try_into().unwrap();
 
-        // Convert HashOutput<32> to [u8; 32]
-        let proof: Vec<[u8; 32]> = proof.into_iter().map(|h| h.0).collect();
+        // If this is a code page and we have a stored HMAC for it, send that instead of the Merkle proof.
+        if section_kind == SectionKind::Code {
+            if let Some(Some(hmac)) = self.stored_code_page_hmacs.get(page_index as usize) {
+                // Build a single-element proof slice containing the HMAC
+                let response = GetPageResponse::new(
+                    (*data).try_into().unwrap(),
+                    is_encrypted,
+                    nonce,
+                    1, // n
+                    1, // t
+                    PageProofKind::Hmac,
+                    &[*hmac],
+                )
+                .serialize();
 
-        // Calculate how many proof elements we can send in one message
+                #[cfg(feature = "debug")]
+                debug!(
+                    "Sending HMAC proof for code page {} instead of Merkle proof",
+                    page_index
+                );
+
+                return self
+                    .transport
+                    .exchange(&apdu_continue(response))
+                    .await
+                    .map_err(VAppEngineError::TransportError);
+            }
+        }
+
+        // Fall back to Merkle proof
+        let proof: Vec<[u8; 32]> = proof.into_iter().map(|h| h.0).collect();
         let t = min(proof.len(), GetPageResponse::max_proof_size()) as u8;
 
         // Create the page response
@@ -217,6 +246,7 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
             nonce,
             proof.len() as u8,
             t,
+            PageProofKind::Merkle,
             &proof[0..t as usize],
         )
         .serialize();
@@ -377,6 +407,32 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
         Ok((status, result))
     }
 
+    async fn process_store_page_proof(
+        &mut self,
+        command: &[u8],
+    ) -> Result<(StatusWord, Vec<u8>), VAppEngineError<E>> {
+        let msg = StorePageProofMessage::deserialize(command)?;
+
+        #[cfg(feature = "debug")]
+        debug!("<- StorePageProofMessage(page = {})", msg.page_index);
+
+        if msg.section_kind != SectionKind::Code {
+            // The app will never ask to store an hmac as proof for non-code pages
+            return Err(VAppEngineError::AccessViolation);
+        }
+        if msg.page_index as usize >= self.stored_code_page_hmacs.len() {
+            self.stored_code_page_hmacs
+                .resize(msg.page_index as usize + 1, None);
+        }
+        self.stored_code_page_hmacs[msg.page_index as usize] = Some(msg.hmac);
+
+        Ok(self
+            .transport
+            .exchange(&apdu_continue(vec![]))
+            .await
+            .map_err(VAppEngineError::TransportError)?)
+    }
+
     async fn process_send_buffer_generic(
         &mut self,
         command: &[u8],
@@ -580,6 +636,7 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> VAppEngine<E> {
                 ClientCommandCode::CommitPage => self.process_commit_page(&result).await?,
                 ClientCommandCode::SendBuffer => self.process_send_buffer(&result).await?,
                 ClientCommandCode::ReceiveBuffer => self.process_receive_buffer(&result).await?,
+                ClientCommandCode::StorePageProof => self.process_store_page_proof(&result).await?,
                 ClientCommandCode::SendBufferContinued
                 | ClientCommandCode::GetPageProofContinued
                 | ClientCommandCode::CommitPageProofContinued => {
@@ -693,6 +750,7 @@ impl<E: std::fmt::Debug + Send + Sync + 'static> GenericVanadiumClient<E> {
             engine_to_client_sender,
             client_to_engine_receiver,
             print_writer,
+            stored_code_page_hmacs: Vec::new(),
         };
 
         // Start the VAppEngine in a task

common/src/client_commands.rs

@@ -62,6 +62,7 @@ pub enum ClientCommandCode {
     SendBuffer = 4,
     SendBufferContinued = 5,
     ReceiveBuffer = 6,
+    StorePageProof = 7,
 }
 
 impl TryFrom<u8> for ClientCommandCode {
@@ -76,12 +77,13 @@ impl TryFrom<u8> for ClientCommandCode {
             4 => Ok(ClientCommandCode::SendBuffer),
             5 => Ok(ClientCommandCode::SendBufferContinued),
             6 => Ok(ClientCommandCode::ReceiveBuffer),
+            7 => Ok(ClientCommandCode::StorePageProof),
             _ => Err("Invalid value for ClientCommandCode"),
         }
     }
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 #[repr(u8)]
 pub enum SectionKind {
     Code = 0,
@@ -154,14 +156,33 @@ impl<'a> Message<'a> for GetPageMessage {
 
 /// Message sent by client in response to the VM's GetPageProofMessage
 /// It contains the page's metadata, and the merkle proof of the page (or part of it)
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(u8)]
+pub enum PageProofKind {
+    Merkle = 0,
+    Hmac = 1,
+}
+
+impl TryFrom<u8> for PageProofKind {
+    type Error = &'static str;
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(PageProofKind::Merkle),
+            1 => Ok(PageProofKind::Hmac),
+            _ => Err("Invalid page proof kind"),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct GetPageResponse<'a> {
     pub page_data: &'a [u8; PAGE_SIZE],
-    pub is_encrypted: bool,    // whether the page is encrypted
-    pub nonce: [u8; 12],       // nonce of the page encryption (all zeros if not encrypted)
-    pub n: u8,                 // number of element in the proof
-    pub t: u8,                 // number of proof elements in this message
-    pub proof: &'a [[u8; 32]], // hashes of the proof
+    pub is_encrypted: bool,
+    pub nonce: [u8; 12],
+    pub n: u8,
+    pub t: u8,
+    pub proof_kind: PageProofKind,
+    pub proof: &'a [[u8; 32]], // Merkle proof elements OR single HMAC (len()==1)
 }
 
 impl<'a> GetPageResponse<'a> {
@@ -172,6 +193,7 @@ impl<'a> GetPageResponse<'a> {
         nonce: [u8; 12],
         n: u8,
         t: u8,
+        proof_kind: PageProofKind,
         proof: &'a [[u8; 32]],
     ) -> Self {
         GetPageResponse {
@@ -180,12 +202,14 @@ impl<'a> GetPageResponse<'a> {
             nonce,
             n,
             t,
+            proof_kind,
             proof,
         }
     }
 
     pub const fn max_proof_size() -> usize {
-        (MAX_APDU_DATA_SIZE - PAGE_SIZE - 1 - 12 - 1 - 1) / 32
+        // page_data | n | t | is_encrypted | nonce(12) | proof_kind | proof elements
+        (MAX_APDU_DATA_SIZE - PAGE_SIZE - 1 - 1 - 1 - 12 - 1) / 32
     }
 }
 
@@ -197,13 +221,14 @@ impl<'a> Message<'a> for GetPageResponse<'a> {
         f(&[self.t]);
         f(&[self.is_encrypted as u8]);
         f(&self.nonce);
+        f(&[self.proof_kind as u8]);
         for p in self.proof {
             f(p);
         }
     }
 
     fn deserialize(data: &'a [u8]) -> Result<Self, MessageDeserializationError> {
-        if data.len() < PAGE_SIZE + 1 + 1 + 1 + 12 {
+        if data.len() < PAGE_SIZE + 1 + 1 + 1 + 12 + 1 {
             return Err(MessageDeserializationError::InvalidDataLength);
         }
         let page_data = data[0..PAGE_SIZE].try_into().unwrap();
@@ -217,13 +242,16 @@ impl<'a> Message<'a> for GetPageResponse<'a> {
         } else {
             [0; 12]
         };
-        let proof_len = data.len() - (PAGE_SIZE + 1 + 1 + 1 + 12);
+        let pk_byte = data[PAGE_SIZE + 15];
+        let proof_kind = PageProofKind::try_from(pk_byte)
+            .map_err(|_| MessageDeserializationError::InvalidDataLength)?;
+        let proof_len = data.len() - (PAGE_SIZE + 1 + 1 + 1 + 12 + 1);
         if proof_len % 32 != 0 {
             return Err(MessageDeserializationError::InvalidDataLength);
         }
         let slice_len = proof_len / 32;
         let proof = unsafe {
-            let ptr = data.as_ptr().add(PAGE_SIZE + 1 + 1 + 1 + 12) as *const [u8; 32];
+            let ptr = data.as_ptr().add(PAGE_SIZE + 1 + 1 + 1 + 12 + 1) as *const [u8; 32];
             core::slice::from_raw_parts(ptr, slice_len)
         };
 
@@ -233,10 +261,63 @@ impl<'a> Message<'a> for GetPageResponse<'a> {
             nonce,
             n,
             t,
+            proof_kind,
             proof,
         })
     }
 }
+/// Message sent by the VM instructing the host to store a compact HMAC proof for a code page.
+#[derive(Debug, Clone)]
+pub struct StorePageProofMessage {
+    pub command_code: ClientCommandCode,
+    pub section_kind: SectionKind,
+    pub page_index: u32,
+    pub hmac: [u8; 32],
+}
+
+impl StorePageProofMessage {
+    #[inline]
+    pub fn new(section_kind: SectionKind, page_index: u32, hmac: [u8; 32]) -> Self {
+        StorePageProofMessage {
+            command_code: ClientCommandCode::StorePageProof,
+            section_kind,
+            page_index,
+            hmac,
+        }
+    }
+}
+
+impl<'a> Message<'a> for StorePageProofMessage {
+    #[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());
+        f(&self.hmac);
+    }
+
+    fn deserialize(data: &'a [u8]) -> Result<Self, MessageDeserializationError> {
+        if data.len() != 1 + 1 + 4 + 32 {
+            return Err(MessageDeserializationError::InvalidDataLength);
+        }
+        let command_code = ClientCommandCode::try_from(data[0])
+            .map_err(|_| MessageDeserializationError::InvalidClientCommandCode)?;
+        if !matches!(command_code, ClientCommandCode::StorePageProof) {
+            return Err(MessageDeserializationError::MismatchingClientCommandCode);
+        }
+        let section_kind = SectionKind::try_from(data[1])
+            .map_err(|_| MessageDeserializationError::InvalidSectionKind)?;
+        let page_index = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+        let mut hmac = [0u8; 32];
+        hmac.copy_from_slice(&data[6..38]);
+        Ok(StorePageProofMessage {
+            command_code,
+            section_kind,
+            page_index,
+            hmac,
+        })
+    }
+}
 
 /// Message sent by the VM to request the rest of the proof, if it didn't fit
 /// in GetPageResponse