From 6e2a89327acb60621520993f2f38868e5291184b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sosth=C3=A8ne=20Gu=C3=A9don?= <sosthene@nitrokey.com>
Date: Thu, 30 Mar 2023 14:05:35 +0200
Subject: [PATCH] Use correct length in advance when loading  files

This reduces the need to shuffle the stack
---
 src/lib.rs   | 8 +++++---
 src/state.rs | 3 +++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0de6770..0e64492 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -907,9 +907,12 @@ impl<'a, T: trussed::Client + AuthClient + trussed::client::Ed255> LoadedAuthent
             _ => &[0x53],
         };
         reply.expand(tag)?;
-        let offset = reply.len();
         match container {
-            Container::KeyHistoryObject => self.get_key_history_object(reply.lend())?,
+            Container::KeyHistoryObject => {
+                let offset = reply.len();
+                self.get_key_history_object(reply.lend())?;
+                reply.prepend_len(offset)?;
+            }
             _ => {
                 if !ContainerStorage(container).load(
                     self.trussed,
@@ -920,7 +923,6 @@ impl<'a, T: trussed::Client + AuthClient + trussed::client::Ed255> LoadedAuthent
                 }
             }
         }
-        reply.prepend_len(offset)?;
 
         Ok(())
     }
diff --git a/src/state.rs b/src/state.rs
index 043da1e..ed2b4f2 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -632,6 +632,7 @@ fn load_if_exists_streaming<const R: usize>(
         Ok(r) => {
             read_len += r.data.len();
             file_len = r.len;
+            buffer.append_len(file_len)?;
             buffer.expand(&r.data)?;
         }
         Err(_) => match try_syscall!(client.entry_metadata(location, path.clone())) {
@@ -752,6 +753,7 @@ impl ContainerStorage {
         }
     }
 
+    // Write the length of the file and write
     pub fn load<const R: usize>(
         self,
         client: &mut impl trussed::Client,
@@ -763,6 +765,7 @@ impl ContainerStorage {
         }
 
         if let Some(data) = self.default() {
+            reply.append_len(data.len())?;
             reply.expand(&data)?;
             Ok(true)
         } else {