Repository: openvm-org/openvm
Pull request #1674: fix: misaligned accesses to sha256
Author: login: arayikhalatyan
Created: 2025-05-21T17:32:22Z | Merged: 2025-05-21T18:55:42Z

Fixes misaligned accesses to sha256. Moved Aligned_buf to openvm platform so both `sha2` and `keccak` can access it


Diff:
diff --git a/crates/toolchain/platform/src/alloc.rs b/crates/toolchain/platform/src/alloc.rs
new file mode 100644
index 0000000000..0af25a3671
--- /dev/null
+++ b/crates/toolchain/platform/src/alloc.rs
@@ -0,0 +1,62 @@
+extern crate alloc;
+
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
+use core::ptr::NonNull;
+
+/// Bytes allocated according to the given Layout
+pub struct AlignedBuf {
+    pub ptr: *mut u8,
+    pub layout: Layout,
+}
+
+impl AlignedBuf {
+    /// Allocate a new buffer whose start address is aligned to `align` bytes.
+    /// *NOTE* if `len` is zero then a creates new `NonNull` that is dangling and 16-byte aligned.
+    pub fn uninit(len: usize, align: usize) -> Self {
+        let layout = Layout::from_size_align(len, align).unwrap();
+        if layout.size() == 0 {
+            return Self {
+                ptr: NonNull::<u128>::dangling().as_ptr() as *mut u8,
+                layout,
+            };
+        }
+        // SAFETY: `len` is nonzero
+        let ptr = unsafe { alloc(layout) };
+        if ptr.is_null() {
+            handle_alloc_error(layout);
+        }
+        AlignedBuf { ptr, layout }
+    }
+
+    /// Allocate a new buffer whose start address is aligned to `align` bytes
+    /// and copy the given data into it.
+    ///
+    /// # Safety
+    /// - `bytes` must not be null
+    /// - `len` should not be zero
+    ///
+    /// See [alloc]. In particular `data` should not be empty.
+    pub unsafe fn new(bytes: *const u8, len: usize, align: usize) -> Self {
+        let buf = Self::uninit(len, align);
+        // SAFETY:
+        // - src and dst are not null
+        // - src and dst are allocated for size
+        // - no alignment requirements on u8
+        // - non-overlapping since ptr is newly allocated
+        unsafe {
+            core::ptr::copy_nonoverlapping(bytes, buf.ptr, len);
+        }
+
+        buf
+    }
+}
+
+impl Drop for AlignedBuf {
+    fn drop(&mut self) {
+        if self.layout.size() != 0 {
+            unsafe {
+                dealloc(self.ptr, self.layout);
+            }
+        }
+    }
+}
diff --git a/crates/toolchain/platform/src/lib.rs b/crates/toolchain/platform/src/lib.rs
index 016211936b..b6a9cfc88e 100644
--- a/crates/toolchain/platform/src/lib.rs
+++ b/crates/toolchain/platform/src/lib.rs
@@ -12,6 +12,9 @@ mod getrandom;
 pub mod heap;
 #[cfg(all(feature = "export-libm", target_os = "zkvm"))]
 mod libm_extern;
+#[cfg(target_os = "zkvm")]
+pub mod alloc;
+
 pub mod memory;
 pub mod print;
 #[cfg(feature = "rust-runtime")]
diff --git a/extensions/keccak256/circuit/src/lib.rs b/extensions/keccak256/circuit/src/lib.rs
index 33e5bbfd5b..6ca057597f 100644
--- a/extensions/keccak256/circuit/src/lib.rs
+++ b/extensions/keccak256/circuit/src/lib.rs
@@ -127,9 +127,9 @@ impl<F: PrimeField32> StepExecutorE1<F> for KeccakVmStep {
 
     fn execute_metered(
         &mut self,
-        state: &mut VmStateMut<GuestMemory, MeteredCtx>,
-        instruction: &Instruction<F>,
-        chip_index: usize,
+        _state: &mut VmStateMut<GuestMemory, MeteredCtx>,
+        _instruction: &Instruction<F>,
+        _chip_index: usize,
     ) -> Result<()> {
         todo!()
     }
diff --git a/extensions/keccak256/guest/src/lib.rs b/extensions/keccak256/guest/src/lib.rs
index 0e5fbfffe9..86ba45e577 100644
--- a/extensions/keccak256/guest/src/lib.rs
+++ b/extensions/keccak256/guest/src/lib.rs
@@ -3,7 +3,7 @@
 #[cfg(target_os = "zkvm")]
 extern crate alloc;
 #[cfg(target_os = "zkvm")]
-use {core::mem::MaybeUninit, zkvm::*};
+use {core::mem::MaybeUninit, openvm_platform::alloc::AlignedBuf};
 
 /// This is custom-0 defined in RISC-V spec document
 pub const OPCODE: u8 = 0x0b;
@@ -27,28 +27,6 @@ pub fn keccak256(input: &[u8]) -> [u8; 32] {
     }
 }
 
-/// keccak256 intrinsic binding
-///
-/// # Safety
-///
-/// The VM accepts the preimage by pointer and length, and writes the
-/// 32-byte hash.
-/// - `bytes` must point to an input buffer at least `len` long.
-/// - `output` must point to a buffer that is at least 32-bytes long.
-/// - `bytes` and `output` must be 4-byte aligned.
-#[cfg(target_os = "zkvm")]
-#[inline(always)]
-unsafe fn __native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
-    openvm_platform::custom_insn_r!(
-        opcode = OPCODE,
-        funct3 = KECCAK256_FUNCT3,
-        funct7 = KECCAK256_FUNCT7,
-        rd = In output,
-        rs1 = In bytes,
-        rs2 = In len
-    );
-}
-
 /// Native hook for keccak256 for use with `alloy-primitives` "native-keccak" feature.
 ///
 /// # Safety
@@ -67,11 +45,12 @@ unsafe fn __native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
 extern "C" fn native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
     // SAFETY: assuming safety assumptions of the inputs, we handle all cases where `bytes` or
     // `output` are not aligned to 4 bytes.
+    const MIN_ALIGN: usize = 4;
     unsafe {
         if bytes as usize % MIN_ALIGN != 0 {
-            let aligned_buff = AlignedBuf::new(bytes, len);
+            let aligned_buff = AlignedBuf::new(bytes, len, MIN_ALIGN);
             if output as usize % MIN_ALIGN != 0 {
-                let aligned_out = AlignedBuf::uninit(32);
+                let aligned_out = AlignedBuf::uninit(32, MIN_ALIGN);
                 __native_keccak256(aligned_buff.ptr, len, aligned_out.ptr);
                 core::ptr::copy_nonoverlapping(aligned_out.ptr as *const u8, output, 32);
             } else {
@@ -79,7 +58,7 @@ extern "C" fn native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
             }
         } else {
             if output as usize % MIN_ALIGN != 0 {
-                let aligned_out = AlignedBuf::uninit(32);
+                let aligned_out = AlignedBuf::uninit(32, MIN_ALIGN);
                 __native_keccak256(bytes, len, aligned_out.ptr);
                 core::ptr::copy_nonoverlapping(aligned_out.ptr as *const u8, output, 32);
             } else {
@@ -89,6 +68,28 @@ extern "C" fn native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
     }
 }
 
+/// keccak256 intrinsic binding
+///
+/// # Safety
+///
+/// The VM accepts the preimage by pointer and length, and writes the
+/// 32-byte hash.
+/// - `bytes` must point to an input buffer at least `len` long.
+/// - `output` must point to a buffer that is at least 32-bytes long.
+/// - `bytes` and `output` must be 4-byte aligned.
+#[cfg(target_os = "zkvm")]
+#[inline(always)]
+fn __native_keccak256(bytes: *const u8, len: usize, output: *mut u8) {
+    openvm_platform::custom_insn_r!(
+        opcode = OPCODE,
+        funct3 = KECCAK256_FUNCT3,
+        funct7 = KECCAK256_FUNCT7,
+        rd = In output,
+        rs1 = In bytes,
+        rs2 = In len
+    );
+}
+
 /// Sets `output` to the keccak256 hash of `input`.
 pub fn set_keccak256(input: &[u8], output: &mut [u8; 32]) {
     #[cfg(not(target_os = "zkvm"))]
@@ -101,70 +102,3 @@ pub fn set_keccak256(input: &[u8], output: &mut [u8; 32]) {
     #[cfg(target_os = "zkvm")]
     native_keccak256(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
 }
-
-#[cfg(target_os = "zkvm")]
-mod zkvm {
-    use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
-    use core::ptr::NonNull;
-
-    use super::*;
-
-    pub const MIN_ALIGN: usize = 4;
-
-    /// Bytes aligned to 4 bytes.
-    pub struct AlignedBuf {
-        pub ptr: *mut u8,
-        pub layout: Layout,
-    }
-
-    impl AlignedBuf {
-        /// Allocate a new buffer whose start address is aligned to 4 bytes.
-        pub fn uninit(len: usize) -> Self {
-            let layout = Layout::from_size_align(len, MIN_ALIGN).unwrap();
-            if layout.size() == 0 {
-                return Self {
-                    ptr: NonNull::<u32>::dangling().as_ptr() as *mut u8,
-                    layout,
-                };
-            }
-            // SAFETY: `len` is nonzero
-            let ptr = unsafe { alloc(layout) };
-            if ptr.is_null() {
-                handle_alloc_error(layout);
-            }
-            AlignedBuf { ptr, layout }
-        }
-
-        /// Allocate a new buffer whose start address is aligned to 4 bytes
-        /// and copy the given data into it.
-        ///
-        /// # Safety
-        /// - `bytes` must not be null
-        /// - `len` should not be zero
-        ///
-        /// See [alloc]. In particular `data` should not be empty.
-        pub unsafe fn new(bytes: *const u8, len: usize) -> Self {
-            let buf = Self::uninit(len);
-            // SAFETY:
-            // - src and dst are not null
-            // - src and dst are allocated for size
-            // - no alignment requirements on u8
-            // - non-overlapping since ptr is newly allocated
-            unsafe {
-                core::ptr::copy_nonoverlapping(bytes, buf.ptr, len);
-            }
-
-            buf
-        }
-    }
-
-    impl Drop for AlignedBuf {
-        fn drop(&mut self) {
-            if self.layout.size() != 0 {
-                unsafe {
-                    dealloc(self.ptr, self.layout);
-                }
-            }
-        }
-    }
-}
diff --git a/extensions/sha256/circuit/src/sha256_chip/mod.rs b/extensions/sha256/circuit/src/sha256_chip/mod.rs
index f9be3fc137..cad58e807c 100644
--- a/extensions/sha256/circuit/src/sha256_chip/mod.rs
+++ b/extensions/sha256/circuit/src/sha256_chip/mod.rs
@@ -2,7 +2,9 @@
 //! variable length inputs read from VM memory.
 
 use openvm_circuit::{
-    arch::{NewVmChipWrapper, Result, StepExecutorE1, VmStateMut},
+    arch::{
+        execution_mode::metered::MeteredCtx, NewVmChipWrapper, Result, StepExecutorE1, VmStateMut,
+    },
     system::memory::online::GuestMemory,
 };
 use openvm_circuit_primitives::{
@@ -70,7 +72,7 @@ impl Sha256VmStep {
 impl<F: PrimeField32> StepExecutorE1<F> for Sha256VmStep {
     fn execute_e1<Ctx>(
         &mut self,
-        state: VmStateMut<GuestMemory, Ctx>,
+        state: &mut VmStateMut<GuestMemory, Ctx>,
         instruction: &Instruction<F>,
     ) -> Result<()> {
         let &Instruction {
@@ -104,6 +106,15 @@ impl<F: PrimeField32> StepExecutorE1<F> for Sha256VmStep {
         memory_write(state.memory, e, dst, hasher.finalize().as_ref());
         Ok(())
     }
+
+    fn execute_metered(
+        &mut self,
+        _state: &mut VmStateMut<GuestMemory, MeteredCtx>,
+        _instruction: &Instruction<F>,
+        _chip_index: usize,
+    ) -> Result<()> {
+        todo!()
+    }
 }
 
 pub fn sha256_solve(input_message: &[u8]) -> [u8; SHA256_WRITE_SIZE] {
diff --git a/extensions/sha256/guest/src/lib.rs b/extensions/sha256/guest/src/lib.rs
index cb34bcd5aa..de7bc6fcb9 100644
--- a/extensions/sha256/guest/src/lib.rs
+++ b/extensions/sha256/guest/src/lib.rs
@@ -1,5 +1,8 @@
 #![no_std]
 
+#[cfg(target_os = "zkvm")]
+use openvm_platform::alloc::AlignedBuf;
+
 /// This is custom-0 defined in RISC-V spec document
 pub const OPCODE: u8 = 0x0b;
 pub const SHA256_FUNCT3: u8 = 0b100;
@@ -13,7 +16,8 @@ pub fn sha256(input: &[u8]) -> [u8; 32] {
     output
 }
 
-/// zkvm native implementation of sha256
+/// Native hook for sha256
+///
 /// # Safety
 ///
 /// The VM accepts the preimage by pointer and length, and writes the
@@ -21,11 +25,53 @@ pub fn sha256(input: &[u8]) -> [u8; 32] {
 /// - `bytes` must point to an input buffer at least `len` long.
 /// - `output` must point to a buffer that is at least 32-bytes long.
 ///
-/// [`sha2-256`]: https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf
+/// [`sha2`]: https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf
 #[cfg(target_os = "zkvm")]
 #[inline(always)]
 #[no_mangle]
 extern "C" fn zkvm_sha256_impl(bytes: *const u8, len: usize, output: *mut u8) {
+    // SAFETY: assuming safety assumptions of the inputs, we handle all cases where `bytes` or
+    // `output` are not aligned to 4 bytes.
+    // The minimum alignment required for the input and output buffers
+    const MIN_ALIGN: usize = 4;
+    // The preferred alignment for the input buffer, since the input is read in chunks of 16 bytes
+    const INPUT_ALIGN: usize = 16;
+    // The preferred alignment for the output buffer, since the output is written in chunks of 32 bytes
+    const OUTPUT_ALIGN: usize = 32;
+    unsafe {
+        if bytes as usize % MIN_ALIGN != 0 {
+            let aligned_buff = AlignedBuf::new(bytes, len, INPUT_ALIGN);
+            if output as usize % MIN_ALIGN != 0 {
+                let aligned_out = AlignedBuf::uninit(32, OUTPUT_ALIGN);
+                __native_sha256(aligned_buff.ptr, len, aligned_out.ptr);
+                core::ptr::copy_nonoverlapping(aligned_out.ptr as *const u8, output, 32);
+            } else {
+                __native_sha256(aligned_buff.ptr, len, output);
+            }
+        } else {
+            if output as usize % MIN_ALIGN != 0 {
+                let aligned_out = AlignedBuf::uninit(32, OUTPUT_ALIGN);
+                __native_sha256(bytes, len, aligned_out.ptr);
+                core::ptr::copy_nonoverlapping(aligned_out.ptr as *const u8, output, 32);
+            } else {
+                __native_sha256(bytes, len, output);
+            }
+        };
+    }
+}
+
+/// sha256 intrinsic binding
+///
+/// # Safety
+///
+/// The VM accepts the preimage by pointer and length, and writes the
+/// 32-byte hash.
+/// - `bytes` must point to an input buffer at least `len` long.
+/// - `output` must point to a buffer that is at least 32-bytes long.
+/// - `bytes` and `output` must be 4-byte aligned.
+#[cfg(target_os = "zkvm")]
+#[inline(always)]
+fn __native_sha256(bytes: *const u8, len: usize, output: *mut u8) {
     openvm_platform::custom_insn_r!(opcode = OPCODE, funct3 = SHA256_FUNCT3, funct7 = SHA256_FUNCT7, rd = In output, rs1 = In bytes, rs2 = In len);
 }
 


Reviews:
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:42:36Z
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:43:29Z
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:43:43Z
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:44:52Z
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:46:13Z
user: login: jonathanpwang
body: lgtm except nits

I think `AlignedBuffer` constructors `uninit, new` should be generalized to not hardcode the min alignment.

In terms of the actual min alignment used, just wanted to check for sha what your thinking was
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:47:30Z
user: login: arayikhalatyan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T18:31:50Z

Review comments:
diff_hunk: @@ -0,0 +1,65 @@
+extern crate alloc;
+use {core::mem::MaybeUninit};
+
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
+use core::ptr::NonNull;
+
+// Register, Memory, and IO address spaces require 4-byte alignment
+pub const MIN_ALIGN: usize = 4;
+
+/// Bytes aligned to 4 bytes.
+pub struct AlignedBuf {
+    pub ptr: *mut u8,
+    pub layout: Layout,
+}
+
+impl AlignedBuf {
+    /// Allocate a new buffer whose start address is aligned to 4 bytes.
+    pub fn uninit(len: usize) -> Self {
path: crates/toolchain/platform/src/aligned_buf.rs
user: login: jonathanpwang
body: let's make this take `align: usize` as an argument to generalize
created_at: 2025-05-21T17:42:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,65 @@
+extern crate alloc;
+use {core::mem::MaybeUninit};
+
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
+use core::ptr::NonNull;
+
+// Register, Memory, and IO address spaces require 4-byte alignment
+pub const MIN_ALIGN: usize = 4;
path: crates/toolchain/platform/src/aligned_buf.rs
user: login: jonathanpwang
body: I think this doesn't need to be here and should be in keccak / sha respectively.

For example, doesn't sha actually need MIN_ALIGN = 8? or you will let the adapter handle 8 <> 4?
created_at: 2025-05-21T17:43:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -12,6 +12,9 @@ mod getrandom;
 pub mod heap;
 #[cfg(all(feature = "export-libm", target_os = "zkvm"))]
 mod libm_extern;
+#[cfg(target_os = "zkvm")]
+pub mod aligned_buf;
path: crates/toolchain/platform/src/lib.rs
user: login: jonathanpwang
body: nit: maybe we call it `alloc`?
created_at: 2025-05-21T17:43:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -8,25 +13,71 @@ pub const SHA256_FUNCT7: u8 = 0x1;
 /// The sha256 cryptographic hash function.
 #[inline(always)]
 pub fn sha256(input: &[u8]) -> [u8; 32] {
-    let mut output = [0u8; 32];
-    set_sha256(input, &mut output);
-    output
+    #[cfg(not(target_os = "zkvm"))]
+    {
+        let mut output = [0u8; 32];
+        set_sha256(input, &mut output);
+        output
+    }
+    #[cfg(target_os = "zkvm")]
+    {
+        let mut output = MaybeUninit::<[u8; 32]>::uninit();
+        zkvm_sha256_impl(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
+        unsafe { output.assume_init() }
+    }
 }
 
-/// zkvm native implementation of sha256
+/// sha256 intrinsic binding
+///
+/// # Safety
+///
+/// The VM accepts the preimage by pointer and length, and writes the
+/// 32-byte hash.
+/// - `bytes` must point to an input buffer at least `len` long.
+/// - `output` must point to a buffer that is at least 32-bytes long.
+/// - `bytes` and `output` must be 4-byte aligned.
+#[cfg(target_os = "zkvm")]
+#[inline(always)]
+fn __native_sha256(bytes: *const u8, len: usize, output: *mut u8) {
path: extensions/sha256/guest/src/lib.rs
user: login: jonathanpwang
body: nit move this private function under the extern function
created_at: 2025-05-21T17:44:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -8,25 +13,71 @@ pub const SHA256_FUNCT7: u8 = 0x1;
 /// The sha256 cryptographic hash function.
 #[inline(always)]
 pub fn sha256(input: &[u8]) -> [u8; 32] {
-    let mut output = [0u8; 32];
-    set_sha256(input, &mut output);
-    output
+    #[cfg(not(target_os = "zkvm"))]
+    {
+        let mut output = [0u8; 32];
+        set_sha256(input, &mut output);
path: extensions/sha256/guest/src/lib.rs
user: login: jonathanpwang
body: to reduce the diff (for security), can we keep `sha256` the same? I think it will amount to the same thing in the end anyways
created_at: 2025-05-21T17:46:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,65 @@
+extern crate alloc;
+use {core::mem::MaybeUninit};
+
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
+use core::ptr::NonNull;
+
+// Register, Memory, and IO address spaces require 4-byte alignment
+pub const MIN_ALIGN: usize = 4;
path: crates/toolchain/platform/src/aligned_buf.rs
user: login: arayikhalatyan
body: That's a good point, changed so when we need to allocate a new memory we allocate with 'preferred' alignment. 
created_at: 2025-05-21T18:31:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: codspeed-hq[bot]
created_at: 2025-05-21T17:36:36Z
body: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_WALLTIME_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Walltime Performance Report](https://codspeed.io/openvm-org/openvm/branches/fix%2Fmisaligned-sha)

### Merging #1674 will **not alter performance**

<sub>Comparing <code>fix/misaligned-sha</code> (f6c8dd2) with <code>feat/new-execution</code> (e8afdad)</sub>


<details>
<summary>:warning: <b>Unknown Walltime execution environment detected</b></summary>

> Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.
>
> For the most accurate results, we recommend using [CodSpeed Macro Runners](https://docs.codspeed.io/instruments/walltime/): bare-metal machines fine-tuned for performance measurement consistency.

</details>

### Summary

`✅ 6` untouched benchmarks  




author_association: NONE
user: login: codspeed-hq[bot]
created_at: 2025-05-21T17:39:23Z
body: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_INSTRUMENTATION_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Instrumentation Performance Report](https://codspeed.io/openvm-org/openvm/branches/fix%2Fmisaligned-sha)

### Merging #1674 will **not alter performance**

<sub>Comparing <code>fix/misaligned-sha</code> (f6c8dd2) with <code>feat/new-execution</code> (e8afdad)</sub>



### Summary

`✅ 6` untouched benchmarks  




author_association: NONE