Merge pull request #34 from senamakel/fast-module-hash

perf(module): hash artifacts with sha2 in large chunks and remember the library digest

Author: senamakel

Cargo.lock

@@ -25,6 +25,7 @@ tokio = { version = "1", features = ["rt-multi-thread", "sync", "time"] }
 async-trait = "0.1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+sha2 = "0.10"
 thiserror = "2"
 tracing = "0.1"
 

Cargo.toml

@@ -24,6 +24,7 @@ autoexamples = false
 async-trait = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+sha2 = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 tokio = { version = "1", features = ["macros", "rt", "sync", "time", "io-util"] }

crates/tinybus/Cargo.toml

@@ -52,6 +52,7 @@ use tracing::{debug, warn};
 
 use crate::error::{Error, Result};
 use crate::module::hash::file_hex;
+use crate::module::remembered_hash::file_hex_remembered;
 
 /// Suffix of the marker beside an archive holding the digest its release
 /// manifest published.
@@ -196,7 +197,7 @@ fn sidecar_matches(module: &Path) -> bool {
     let Ok(file) = std::fs::File::open(module) else {
         return false;
     };
-    file_hex(file).is_ok_and(|actual| actual == expected)
+    file_hex_remembered(file).is_ok_and(|actual| actual == expected)
 }
 
 /// Where the digest marker for `asset_name` lives in `dir`.

crates/tinybus/src/module/cache.rs

@@ -1,102 +1,43 @@
-//! Dependency-free SHA-256 used by module artifact verification.
+//! SHA-256 used by module artifact verification.
 //!
 //! Verification at load is what an attestation later stands on: a recipient
 //! earns the right to be handed a confidential body by having its artifact
 //! digested here and matched against the operator's allowlist. Compiled
 //! unconditionally rather than behind `modules`, because `module::sha256_file`
 //! offers the same digest to callers that publish an asset without ever
 //! loading one.
+//!
+//! Input is read in large chunks, not one 64-byte block per `read` syscall, and
+//! the compression function is `sha2`'s, which uses the CPU's SHA extensions
+//! where it has them, so hashing a 25 MB library stays off the startup path.
+//! Hashing the same library twice is avoided one level up, in
+//! `remembered_hash`.
 
 use std::io::{self, Read};
 
-const INITIAL: [u32; 8] = [
-    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
-];
+use sha2::{Digest, Sha256};
 
-const K: [u32; 64] = [
-    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
-    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
-    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
-    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
-    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
-    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
-    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
-    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
-];
+/// Bytes requested per `read`. Large enough that a 25 MB library is ~100
+/// syscalls, small enough to stay out of the allocator's way.
+const CHUNK: usize = 256 * 1024;
 
+/// Lowercase hex SHA-256 of everything `reader` yields.
 pub(crate) fn file_hex(mut reader: impl Read) -> io::Result<String> {
-    let mut state = INITIAL;
-    let mut block = [0u8; 64];
-    let mut used = 0usize;
-    let mut byte_len = 0u64;
+    let mut hasher = Sha256::new();
+    let mut chunk = vec![0u8; CHUNK];
     loop {
-        let read = reader.read(&mut block[used..])?;
-        if read == 0 {
-            break;
-        }
-        used += read;
-        byte_len = byte_len.wrapping_add(read as u64);
-        if used == block.len() {
-            compress(&mut state, &block);
-            used = 0;
+        match reader.read(&mut chunk) {
+            Ok(0) => break,
+            Ok(read) => hasher.update(&chunk[..read]),
+            Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
+            Err(error) => return Err(error),
         }
     }
-    block[used] = 0x80;
-    used += 1;
-    if used > 56 {
-        block[used..].fill(0);
-        compress(&mut state, &block);
-        block.fill(0);
-    } else {
-        block[used..56].fill(0);
-    }
-    block[56..].copy_from_slice(&byte_len.wrapping_mul(8).to_be_bytes());
-    compress(&mut state, &block);
-    Ok(state.iter().map(|word| format!("{word:08x}")).collect())
-}
-
-fn compress(state: &mut [u32; 8], block: &[u8; 64]) {
-    let mut words = [0u32; 64];
-    for (index, chunk) in block.as_chunks::<4>().0.iter().enumerate() {
-        words[index] = u32::from_be_bytes(*chunk);
-    }
-    for index in 16..64 {
-        let s0 = words[index - 15].rotate_right(7)
-            ^ words[index - 15].rotate_right(18)
-            ^ (words[index - 15] >> 3);
-        let s1 = words[index - 2].rotate_right(17)
-            ^ words[index - 2].rotate_right(19)
-            ^ (words[index - 2] >> 10);
-        words[index] = words[index - 16]
-            .wrapping_add(s0)
-            .wrapping_add(words[index - 7])
-            .wrapping_add(s1);
-    }
-
-    let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state;
-    for index in 0..64 {
-        let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
-        let choice = (e & f) ^ ((!e) & g);
-        let t1 = h
-            .wrapping_add(s1)
-            .wrapping_add(choice)
-            .wrapping_add(K[index])
-            .wrapping_add(words[index]);
-        let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
-        let majority = (a & b) ^ (a & c) ^ (b & c);
-        let t2 = s0.wrapping_add(majority);
-        h = g;
-        g = f;
-        f = e;
-        e = d.wrapping_add(t1);
-        d = c;
-        c = b;
-        b = a;
-        a = t1.wrapping_add(t2);
-    }
-    for (slot, value) in state.iter_mut().zip([a, b, c, d, e, f, g, h]) {
-        *slot = slot.wrapping_add(value);
-    }
+    Ok(hasher
+        .finalize()
+        .iter()
+        .map(|byte| format!("{byte:02x}"))
+        .collect())
 }
 
 #[cfg(test)]

crates/tinybus/src/module/hash.rs

@@ -21,3 +21,83 @@ fn sha256_matches_a_published_multi_block_input_vector() {
         "41edece42d63e8d9bf515a9ba6932e1c20cbc9f5a5d134645adb5db1b9737ea3"
     );
 }
+
+/// Hands out at most `step` bytes per `read`, the way a pipe or a slow disk can.
+struct Trickle<'a> {
+    data: &'a [u8],
+    step: usize,
+}
+
+impl std::io::Read for Trickle<'_> {
+    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
+        let n = self.step.min(out.len()).min(self.data.len());
+        out[..n].copy_from_slice(&self.data[..n]);
+        self.data = &self.data[n..];
+        Ok(n)
+    }
+}
+
+fn pattern(len: usize) -> Vec<u8> {
+    (0..len).map(|i| (i * 31 % 251) as u8).collect()
+}
+
+#[test]
+fn sha256_is_the_same_whatever_the_read_sizes_and_chunk_boundaries() {
+    // Spans several chunks and ends mid-chunk and mid-block.
+    let data = pattern(super::CHUNK * 3 + 77);
+    let whole = super::file_hex(&data[..]).unwrap();
+    for step in [1, 63, 64, 65, 4096, super::CHUNK - 1, super::CHUNK + 1] {
+        let trickled = super::file_hex(Trickle { data: &data, step }).unwrap();
+        assert_eq!(trickled, whole, "step {step}");
+    }
+}
+
+#[test]
+fn sha256_matches_the_published_million_a_vector() {
+    assert_eq!(
+        super::file_hex(&vec![b'a'; 1_000_000][..]).unwrap(),
+        "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0"
+    );
+}
+
+/// Fails with `kind` on the first `read`, then behaves like `data`.
+struct FailsOnce<'a> {
+    data: &'a [u8],
+    kind: Option<std::io::ErrorKind>,
+}
+
+impl std::io::Read for FailsOnce<'_> {
+    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
+        if let Some(kind) = self.kind.take() {
+            return Err(kind.into());
+        }
+        let n = out.len().min(self.data.len());
+        out[..n].copy_from_slice(&self.data[..n]);
+        self.data = &self.data[n..];
+        Ok(n)
+    }
+}
+
+#[test]
+fn an_interrupted_read_is_retried_not_reported() {
+    let data = pattern(5000);
+    let reader = FailsOnce {
+        data: &data,
+        kind: Some(std::io::ErrorKind::Interrupted),
+    };
+    assert_eq!(
+        super::file_hex(reader).unwrap(),
+        super::file_hex(&data[..]).unwrap()
+    );
+}
+
+#[test]
+fn a_failing_read_is_an_error_not_a_digest_of_what_was_read() {
+    let data = pattern(5000);
+    let reader = FailsOnce {
+        data: &data,
+        kind: Some(std::io::ErrorKind::PermissionDenied),
+    };
+    let error = super::file_hex(reader).unwrap_err();
+    assert_eq!(error.kind(), std::io::ErrorKind::PermissionDenied);
+}

crates/tinybus/src/module/hash_tests.rs

@@ -1842,7 +1842,7 @@ fn allowlisted_hash(path: &Path, file: std::fs::File) -> Result<Option<String>>
             "module allowlist contains an invalid hash",
         ));
     }
-    let actual = crate::module::hash::file_hex(file)
+    let actual = crate::module::remembered_hash::file_hex_remembered(file)
         .map_err(|_| Error::module_refused(path, "artifact hash could not be read"))?;
     if actual != expected {
         return Err(Error::module_refused(

crates/tinybus/src/module/host.rs

@@ -22,6 +22,8 @@ pub(crate) mod host;
 #[cfg(feature = "modules")]
 mod loader;
 #[cfg(feature = "modules")]
+mod remembered_hash;
+#[cfg(feature = "modules")]
 mod resolve;
 #[cfg(feature = "modules")]
 mod transport;