RR #2: Sha256 checksum for components (#12576)

* Add sha256 checksum for component for record/replay consistency

* Move sha2 crate as workspace dependency

* Run checksum digest only on recording configs

* Fix CI error and restructure from_binary

Author: arjunr2

Cargo.lock

@@ -437,6 +437,7 @@ arbtest = "0.3.2"
 rayon = "1.5.3"
 regex = "1.9.1"
 pin-project-lite = "0.2.14"
+sha2 = { version = "0.10.2", default-features = false }
 
 # =============================================================================
 #

Cargo.toml

@@ -19,7 +19,7 @@ directories-next = "2.0"
 log = { workspace = true }
 serde = { workspace = true }
 serde_derive = { workspace = true }
-sha2 = "0.10.2"
+sha2 = { workspace = true, features = ['std'] }
 toml = { workspace = true }
 zstd = { version = "0.13.0", default-features = false }
 wasmtime-environ = { workspace = true }

crates/cache/Cargo.toml

@@ -39,6 +39,7 @@ wasmtime-component-util = { workspace = true, optional = true }
 wasmtime-core = { workspace = true }
 semver = { workspace = true, optional = true, features = ['serde'] }
 smallvec = { workspace = true, features = ['serde'] }
+sha2 = { workspace = true }
 
 [dev-dependencies]
 clap = { workspace = true, features = ['default'] }
@@ -85,3 +86,6 @@ std = [
   'indexmap/std',
   'wasmtime-core/std',
 ]
+rr = [
+  "component-model"
+]

crates/environ/Cargo.toml

@@ -1,6 +1,7 @@
 //! Definitions of runtime structures and metadata which are serialized into ELF
 //! with `postcard` as part of a module's compilation process.
 
+use crate::WasmChecksum;
 use crate::error::{Result, bail};
 use crate::prelude::*;
 use crate::{
@@ -120,6 +121,7 @@ impl<'a> ObjectBuilder<'a> {
             data,
             data_align,
             passive_data,
+            wasm,
             ..
         } = translation;
 
@@ -220,6 +222,7 @@ impl<'a> ObjectBuilder<'a> {
                 has_wasm_debuginfo: self.tunables.parse_wasm_debuginfo,
                 dwarf,
             },
+            checksum: WasmChecksum::from_binary(wasm, self.tunables.recording),
         })
     }
 

crates/environ/src/compile/module_artifacts.rs

@@ -2,7 +2,7 @@
 //! which are serialized with `bincode` into output ELF files.
 
 use crate::{
-    CompiledFunctionsTable, CompiledModuleInfo, PrimaryMap, StaticModuleIndex,
+    CompiledFunctionsTable, CompiledModuleInfo, PrimaryMap, StaticModuleIndex, WasmChecksum,
     component::{Component, ComponentTypes, TypeComponentIndex},
 };
 use serde_derive::{Deserialize, Serialize};
@@ -20,6 +20,8 @@ pub struct ComponentArtifacts {
     pub types: ComponentTypes,
     /// Serialized metadata about all included core wasm modules.
     pub static_modules: PrimaryMap<StaticModuleIndex, CompiledModuleInfo>,
+    /// A checksum of the source Wasm binary from which the component was compiled.
+    pub checksum: WasmChecksum,
 }
 
 /// Runtime state that a component retains to support its operation.

crates/environ/src/component/artifacts.rs

@@ -8,6 +8,8 @@ use core::{fmt, u32};
 use core::{iter, str};
 use cranelift_entity::{EntityRef, PrimaryMap};
 use serde_derive::{Deserialize, Serialize};
+#[cfg(feature = "rr")]
+use sha2::{Digest, Sha256};
 
 /// Description of where a function is located in the text section of a
 /// compiled image.
@@ -28,6 +30,42 @@ impl FunctionLoc {
     }
 }
 
+/// The checksum of a Wasm binary.
+///
+/// Allows for features requiring the exact same Wasm Module (e.g. deterministic replay)
+/// to verify that the binary used matches the one originally compiled.
+#[derive(Copy, Clone, Default, PartialEq, Eq, Ord, PartialOrd, Debug, Serialize, Deserialize)]
+pub struct WasmChecksum([u8; 32]);
+
+impl WasmChecksum {
+    /// Construct a [`WasmChecksum`] from the given wasm binary, used primarily for integrity
+    /// checks on compiled modules when recording configs are enabled. The checksum is not
+    /// computed when recording is disabled to prevent pessimization of non-recorded compilations.
+    #[cfg(feature = "rr")]
+    pub fn from_binary(bin: &[u8], recording: bool) -> WasmChecksum {
+        if recording {
+            WasmChecksum(Sha256::digest(bin).into())
+        } else {
+            WasmChecksum::default()
+        }
+    }
+
+    /// This method requires the `rr` feature to actual compute a checksum. Since the `rr`
+    /// feature is disabled, this only returns a default checksum value of all zeros.
+    #[cfg(not(feature = "rr"))]
+    pub fn from_binary(_: &[u8], _: bool) -> WasmChecksum {
+        WasmChecksum::default()
+    }
+}
+
+impl core::ops::Deref for WasmChecksum {
+    type Target = [u8; 32];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
 /// A builder for a `CompiledFunctionsTable`.
 pub struct CompiledFunctionsTableBuilder {
     inner: CompiledFunctionsTable,
@@ -548,6 +586,9 @@ pub struct CompiledModuleInfo {
 
     /// Sorted list, by function index, of names we have for this module.
     pub func_names: Vec<FunctionName>,
+
+    /// Checksum of the source Wasm binary from which this module was compiled.
+    pub checksum: WasmChecksum,
 }
 
 /// The name of a function stored in the

crates/environ/src/module_artifacts.rs

@@ -144,6 +144,10 @@ define_tunables! {
         /// Whether any component model feature related to concurrency is
         /// enabled.
         pub concurrency_support: bool,
+
+        /// Whether recording in RR is enabled or not. This is used primarily
+        /// to signal checksum computation for compiled artifacts.
+        pub recording: bool,
     }
 
     pub struct ConfigTunables {
@@ -220,6 +224,7 @@ impl Tunables {
             inlining_sum_size_threshold: 2000,
             debug_guest: false,
             concurrency_support: true,
+            recording: false,
         }
     }
 

crates/environ/src/tunables.rs

@@ -17,7 +17,7 @@ wit-bindgen = { workspace = true, features = ['default', 'async-spawn', 'inter-t
 libc = { workspace = true }
 futures = { workspace = true, default-features = false, features = ['alloc', 'async-await'] }
 url = { workspace = true }
-sha2 = "0.10.2"
+sha2 = { workspace = true, features = ['std'] }
 base64 = { workspace = true }
 wasip1 = { version = "1.0.0", default-features = true }
 wasip2 = "1.0.0"

crates/test-programs/Cargo.toml

@@ -50,7 +50,7 @@ tracing-subscriber = { workspace = true }
 wasmtime = { workspace = true, features = ['default', 'anyhow'] }
 tokio = { workspace = true, features = ['fs', 'macros'] }
 futures = { workspace = true, default-features = false, features = ['alloc', 'async-await'] }
-sha2 = "0.10.2"
+sha2 = { workspace = true, features = ['std'] }
 base64 = { workspace = true }
 flate2 = { workspace = true }
 wasm-compose = { workspace = true }