6 files changed, 140 insertions, 37 deletions

diff --git a/embassy-boot/boot/Cargo.toml b/embassy-boot/boot/Cargo.toml
index 3312c2f9f..c4ebdaff2 100644
--- a/embassy-boot/boot/Cargo.toml
+++ b/embassy-boot/boot/Cargo.toml
@@ -24,6 +24,7 @@ features = ["defmt"]
 
 [dependencies]
 defmt = { version = "0.3", optional = true }
+digest = "0.10"
 log = { version = "0.4", optional = true  }
 ed25519-dalek = { version = "1.0.1", default_features = false, features = ["u32_backend"], optional = true }
 embassy-sync = { version = "0.1.0", path = "../../embassy-sync" }
@@ -37,6 +38,7 @@ log = "0.4"
 env_logger = "0.9"
 rand = "0.7" # ed25519-dalek v1.0.1 depends on this exact version
 futures = { version = "0.3", features = ["executor"] }
+sha1 = "0.10.5"
 
 [dev-dependencies.ed25519-dalek]
 default_features = false
@@ -47,4 +49,4 @@ ed25519-dalek = ["dep:ed25519-dalek", "_verify"]
 ed25519-salty = ["dep:salty", "_verify"]
 
 #Internal features
-_verify = []
\ No newline at end of file
+_verify = []
diff --git a/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
new file mode 100644
index 000000000..a184d1c51
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
@@ -0,0 +1,30 @@
+use digest::typenum::U64;
+use digest::{FixedOutput, HashMarker, OutputSizeUser, Update};
+use ed25519_dalek::Digest as _;
+
+pub struct Sha512(ed25519_dalek::Sha512);
+
+impl Default for Sha512 {
+    fn default() -> Self {
+        Self(ed25519_dalek::Sha512::new())
+    }
+}
+
+impl Update for Sha512 {
+    fn update(&mut self, data: &[u8]) {
+        self.0.update(data)
+    }
+}
+
+impl FixedOutput for Sha512 {
+    fn finalize_into(self, out: &mut digest::Output<Self>) {
+        let result = self.0.finalize();
+        out.as_mut_slice().copy_from_slice(result.as_slice())
+    }
+}
+
+impl OutputSizeUser for Sha512 {
+    type OutputSize = U64;
+}
+
+impl HashMarker for Sha512 {}
diff --git a/embassy-boot/boot/src/digest_adapters/mod.rs b/embassy-boot/boot/src/digest_adapters/mod.rs
new file mode 100644
index 000000000..9b4b4b60c
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/mod.rs
@@ -0,0 +1,5 @@
+#[cfg(feature = "ed25519-dalek")]
+pub(crate) mod ed25519_dalek;
+
+#[cfg(feature = "ed25519-salty")]
+pub(crate) mod salty;
diff --git a/embassy-boot/boot/src/digest_adapters/salty.rs b/embassy-boot/boot/src/digest_adapters/salty.rs
new file mode 100644
index 000000000..2b5dcf3af
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/salty.rs
@@ -0,0 +1,29 @@
+use digest::typenum::U64;
+use digest::{FixedOutput, HashMarker, OutputSizeUser, Update};
+
+pub struct Sha512(salty::Sha512);
+
+impl Default for Sha512 {
+    fn default() -> Self {
+        Self(salty::Sha512::new())
+    }
+}
+
+impl Update for Sha512 {
+    fn update(&mut self, data: &[u8]) {
+        self.0.update(data)
+    }
+}
+
+impl FixedOutput for Sha512 {
+    fn finalize_into(self, out: &mut digest::Output<Self>) {
+        let result = self.0.finalize();
+        out.as_mut_slice().copy_from_slice(result.as_slice())
+    }
+}
+
+impl OutputSizeUser for Sha512 {
+    type OutputSize = U64;
+}
+
+impl HashMarker for Sha512 {}
diff --git a/embassy-boot/boot/src/firmware_updater.rs b/embassy-boot/boot/src/firmware_updater.rs
index 22e3e6b00..2d1b26980 100644
--- a/embassy-boot/boot/src/firmware_updater.rs
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -1,3 +1,4 @@
+use digest::Digest;
 use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind};
 use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
 
@@ -128,25 +129,27 @@ impl FirmwareUpdater {
 
         #[cfg(feature = "ed25519-dalek")]
         {
-            use ed25519_dalek::{Digest, PublicKey, Sha512, Signature, SignatureError, Verifier};
+            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+
+            use crate::digest_adapters::ed25519_dalek::Sha512;
 
             let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
 
             let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
-            let mut digest = Sha512::new();
-            self.incremental_hash(_state_and_dfu_flash, _update_len, _aligned, |x| digest.update(x))
+            let mut message = [0; 64];
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
                 .await?;
 
-            public_key
-                .verify(&digest.finalize(), &signature)
-                .map_err(into_signature_error)?
+            public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
         #[cfg(feature = "ed25519-salty")]
         {
             use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
-            use salty::{PublicKey, Sha512, Signature};
+            use salty::{PublicKey, Signature};
+
+            use crate::digest_adapters::salty::Sha512;
 
             fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
                 FirmwareUpdaterError::Signature(signature::Error::default())
@@ -157,11 +160,10 @@ impl FirmwareUpdater {
             let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
-            let mut digest = Sha512::new();
-            self.incremental_hash(_state_and_dfu_flash, _update_len, _aligned, |x| digest.update(x))
+            let mut message = [0; 64];
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
                 .await?;
 
-            let message = digest.finalize();
             let r = public_key.verify(&message, &signature);
             trace!(
                 "Verifying with public key {}, signature {} and message {} yields ok: {}",
@@ -176,19 +178,21 @@ impl FirmwareUpdater {
         self.set_magic(_aligned, SWAP_MAGIC, _state_and_dfu_flash).await
     }
 
-    /// Iterate through the DFU and process all bytes with the provided closure.
-    pub async fn incremental_hash<F: AsyncNorFlash>(
+    /// Verify the update in DFU with any digest.
+    pub async fn hash<F: AsyncNorFlash, D: Digest>(
         &mut self,
         dfu_flash: &mut F,
         update_len: u32,
-        aligned: &mut [u8],
-        mut update: impl FnMut(&[u8]),
+        chunk_buf: &mut [u8],
+        output: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
-        for offset in (0..update_len).step_by(aligned.len()) {
-            self.dfu.read(dfu_flash, offset, aligned).await?;
-            let len = core::cmp::min((update_len - offset) as usize, aligned.len());
-            update(&aligned[..len]);
+        let mut digest = D::new();
+        for offset in (0..update_len).step_by(chunk_buf.len()) {
+            self.dfu.read(dfu_flash, offset, chunk_buf).await?;
+            let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
+            digest.update(&chunk_buf[..len]);
         }
+        output.copy_from_slice(digest.finalize().as_slice());
         Ok(())
     }
 
@@ -334,24 +338,26 @@ impl FirmwareUpdater {
 
         #[cfg(feature = "ed25519-dalek")]
         {
-            use ed25519_dalek::{Digest, PublicKey, Sha512, Signature, SignatureError, Verifier};
+            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+
+            use crate::digest_adapters::ed25519_dalek::Sha512;
 
             let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
 
             let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
-            let mut digest = Sha512::new();
-            self.incremental_hash_blocking(_state_and_dfu_flash, _update_len, _aligned, |x| digest.update(x))?;
+            let mut message = [0; 64];
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
 
-            public_key
-                .verify(&digest.finalize(), &signature)
-                .map_err(into_signature_error)?
+            public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
         #[cfg(feature = "ed25519-salty")]
         {
             use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
-            use salty::{PublicKey, Sha512, Signature};
+            use salty::{PublicKey, Signature};
+
+            use crate::digest_adapters::salty::Sha512;
 
             fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
                 FirmwareUpdaterError::Signature(signature::Error::default())
@@ -362,10 +368,9 @@ impl FirmwareUpdater {
             let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
-            let mut digest = Sha512::new();
-            self.incremental_hash_blocking(_state_and_dfu_flash, _update_len, _aligned, |x| digest.update(x))?;
+            let mut message = [0; 64];
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
 
-            let message = digest.finalize();
             let r = public_key.verify(&message, &signature);
             trace!(
                 "Verifying with public key {}, signature {} and message {} yields ok: {}",
@@ -380,19 +385,21 @@ impl FirmwareUpdater {
         self.set_magic_blocking(_aligned, SWAP_MAGIC, _state_and_dfu_flash)
     }
 
-    /// Iterate through the DFU and process all bytes with the provided closure.
-    pub fn incremental_hash_blocking<F: NorFlash>(
+    /// Verify the update in DFU with any digest.
+    pub fn hash_blocking<F: NorFlash, D: Digest>(
         &mut self,
         dfu_flash: &mut F,
         update_len: u32,
-        aligned: &mut [u8],
-        mut update: impl FnMut(&[u8]),
+        chunk_buf: &mut [u8],
+        output: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
-        for offset in (0..update_len).step_by(aligned.len()) {
-            self.dfu.read_blocking(dfu_flash, offset, aligned)?;
-            let len = core::cmp::min((update_len - offset) as usize, aligned.len());
-            update(&aligned[..len]);
+        let mut digest = D::new();
+        for offset in (0..update_len).step_by(chunk_buf.len()) {
+            self.dfu.read_blocking(dfu_flash, offset, chunk_buf)?;
+            let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
+            digest.update(&chunk_buf[..len]);
         }
+        output.copy_from_slice(digest.finalize().as_slice());
         Ok(())
     }
 
@@ -479,3 +486,32 @@ impl FirmwareUpdater {
         Ok(self.dfu)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use futures::executor::block_on;
+    use sha1::{Digest, Sha1};
+
+    use super::*;
+    use crate::tests::MemFlash;
+
+    #[test]
+    fn can_verify() {
+        const STATE: Partition = Partition::new(0, 4096);
+        const DFU: Partition = Partition::new(65536, 131072);
+
+        let mut flash = MemFlash::<131072, 4096, 8>([0xFF; 131072]);
+
+        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
+        let mut to_write = [0; 4096];
+        to_write[..7].copy_from_slice(update.as_slice());
+
+        let mut updater = FirmwareUpdater::new(DFU, STATE);
+        block_on(updater.write_firmware(0, to_write.as_slice(), &mut flash)).unwrap();
+        let mut chunk_buf = [0; 2];
+        let mut hash = [0; 20];
+        block_on(updater.hash::<_, Sha1>(&mut flash, update.len() as u32, &mut chunk_buf, &mut hash)).unwrap();
+
+        assert_eq!(Sha1::digest(update).as_slice(), hash);
+    }
+}
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 6888a8055..da9055476 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -6,6 +6,7 @@
 mod fmt;
 
 mod boot_loader;
+mod digest_adapters;
 mod firmware_updater;
 mod partition;