feat: make keccak/sha256 memory access aligned to 4-byte (#1859)

picked from 1.4 branch.

i think at least this PR will not make anything worse?

Author: lispc

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);
+            }
+        }
+    }
+}

crates/toolchain/platform/src/lib.rs

@@ -4,12 +4,15 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
-#[cfg(all(feature = "rust-runtime", target_os = "zkvm"))]
+#[cfg(target_os = "zkvm")]
 pub use openvm_custom_insn::{custom_insn_i, custom_insn_r};
+#[cfg(target_os = "zkvm")]
+pub mod alloc;
 #[cfg(all(feature = "rust-runtime", target_os = "zkvm"))]
 pub mod heap;
 #[cfg(all(feature = "export-libm", target_os = "zkvm"))]
 mod libm_extern;
+
 pub mod memory;
 pub mod print;
 #[cfg(feature = "rust-runtime")]
@@ -19,9 +22,6 @@ pub mod rust_rt;
 /// 4 bytes (i.e. 32 bits) as the zkVM is an implementation of the rv32im ISA.
 pub const WORD_SIZE: usize = core::mem::size_of::<u32>();
 
-/// Size of a zkVM memory page.
-pub const PAGE_SIZE: usize = 1024;
-
 /// Standard IO file descriptors for use with sys_read and sys_write.
 pub mod fileno {
     pub const STDIN: u32 = 0;

extensions/keccak256/guest/src/lib.rs

@@ -1,5 +1,10 @@
 #![no_std]
 
+#[cfg(target_os = "zkvm")]
+extern crate alloc;
+#[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 KECCAK256_FUNCT3: u8 = 0b100;
@@ -21,6 +26,43 @@ pub const KECCAK256_FUNCT7: u8 = 0;
 #[inline(always)]
 #[no_mangle]
 pub 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, MIN_ALIGN);
+            if output as usize % MIN_ALIGN != 0 {
+                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 {
+                __native_keccak256(aligned_buff.ptr, len, output);
+            }
+        } else {
+            if output as usize % MIN_ALIGN != 0 {
+                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 {
+                __native_keccak256(bytes, len, output);
+            }
+        };
+    }
+}
+
+/// 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,

extensions/sha256/guest/src/lib.rs

@@ -1,22 +1,69 @@
 #![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;
 pub const SHA256_FUNCT7: u8 = 0x1;
 
-/// zkvm native implementation of sha256
+/// Native hook for sha256
+///
 /// # 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.
 ///
-/// [`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]
 pub 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);
 }