Implement Memory Sanitizer unpoisoning more precisely.

.github/workflows/tests.yml

@@ -139,35 +139,65 @@ jobs:
       - run: cargo test --target=x86_64-win7-windows-msvc -Z build-std --features=std
       - run: cargo test --target=i686-win7-windows-msvc -Z build-std --features=std
 
-  sanitizer-linux-aarch64:
-    name: Sanitizer Linux AArch64
-    # MemorySanitizer won't run in QEMU so we can't run it in cross:
-    # https://github.com/llvm/llvm-project/issues/65144
-    runs-on: ubuntu-24.04-arm
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@master
-        with:
-          toolchain: nightly-2025-06-01
-          components: rust-src
-      - env:
-          RUSTFLAGS: -Dwarnings -Zsanitizer=memory
-          RUSTDOCFLAGS: -Dwarnings -Zsanitizer=memory
-        run: cargo test -Zbuild-std --target=aarch64-unknown-linux-gnu
+  sanitizer-linux:
+    name: Sanitizer Linux
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      matrix:
+        arch: [
+            "aarch64",
+            "x86_64",
+        ]
+        include:
+          # MemorySanitizer won't run in QEMU so we can't run it in cross:
+          # https://github.com/llvm/llvm-project/issues/65144
+          - arch: aarch64
+            runner: ubuntu-24.04-arm
 
-  sanitizer-linux-x86_64:
-    name: Sanitizer Linux x86_64
-    runs-on: ubuntu-24.04
+          - arch: x86_64
+            runner: ubuntu-24.04
     steps:
       - uses: actions/checkout@v4
       - uses: dtolnay/rust-toolchain@master
         with:
           toolchain: nightly-2025-06-01
           components: rust-src
-      - env:
+      - name: default configuration
+        env:
           RUSTFLAGS: -Dwarnings -Zsanitizer=memory
           RUSTDOCFLAGS: -Dwarnings -Zsanitizer=memory
-        run: cargo test -Zbuild-std --target=x86_64-unknown-linux-gnu
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - name: --cfg getrandom_backend="linux_getrandom"
+        env:
+          RUSTFLAGS: --cfg getrandom_backend="linux_getrandom" -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_backend="linux_getrandom" -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - name: --cfg getrandom_backend="linux_raw"
+        env:
+          RUSTFLAGS: --cfg getrandom_backend="linux_raw" -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_backend="linux_raw" -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - name: --cfg getrandom_backend="linux_fallback"
+        env:
+          RUSTFLAGS: --cfg getrandom_backend="linux_fallback" -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_backend="linux_fallback" -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - if: ${{ matrix.arch == 'x86_64' }}
+        name: --cfg getrandom_backend="rdrand"
+        env:
+          RUSTFLAGS: --cfg getrandom_backend="rdrand" -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_backend="rdrand" -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - name: --cfg getrandom_test_linux_fallback
+        env:
+          RUSTFLAGS: --cfg getrandom_test_linux_fallback -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_test_linux_fallback -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
+      - name: --cfg getrandom_test_linux_without_fallback
+        env:
+          RUSTFLAGS: --cfg getrandom_test_linux_without_fallback -Dwarnings -Zsanitizer=memory
+          RUSTDOCFLAGS: --cfg getrandom_test_linux_without_fallback -Dwarnings -Zsanitizer=memory
+        run: cargo test -Zbuild-std --target=${{ matrix.arch }}-unknown-linux-gnu
 
   cross:
     name: Cross

src/backends.rs

@@ -2,7 +2,8 @@
 //!
 //! This module should provide `fill_inner` with the signature
 //! `fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error>`.
-//! The function MUST fully initialize `dest` when `Ok(())` is returned.
+//! The function MUST fully initialize `dest` when `Ok(())` is returned;
+//! the function may need to use `sanitizer::unpoison` as well.
 //! The function MUST NOT ever write uninitialized bytes into `dest`,
 //! regardless of what value it returns.
 
@@ -12,9 +13,11 @@ cfg_if! {
         pub use custom::*;
     } else if #[cfg(getrandom_backend = "linux_getrandom")] {
         mod getrandom;
+        mod sanitizer;
         pub use getrandom::*;
     } else if #[cfg(getrandom_backend = "linux_raw")] {
         mod linux_raw;
+        mod sanitizer;
         pub use linux_raw::*;
     } else if #[cfg(getrandom_backend = "rdrand")] {
         mod rdrand;
@@ -43,6 +46,7 @@ cfg_if! {
         pub use unsupported::*;
     } else if #[cfg(all(target_os = "linux", target_env = ""))] {
         mod linux_raw;
+        mod sanitizer;
         pub use linux_raw::*;
     } else if #[cfg(target_os = "espidf")] {
         mod esp_idf;
@@ -102,6 +106,7 @@ cfg_if! {
     ))] {
         mod use_file;
         mod linux_android_with_fallback;
+        mod sanitizer;
         pub use linux_android_with_fallback::*;
     } else if #[cfg(any(
         target_os = "android",
@@ -116,6 +121,8 @@ cfg_if! {
         all(target_os = "horizon", target_arch = "arm"),
     ))] {
         mod getrandom;
+        #[cfg(any(target_os = "android", target_os = "linux"))]
+        mod sanitizer;
         pub use getrandom::*;
     } else if #[cfg(target_os = "solaris")] {
         mod solaris;

src/backends/getrandom.rs

@@ -26,6 +26,14 @@ mod util_libc;
 #[inline]
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
     util_libc::sys_fill_exact(dest, |buf| unsafe {
-        libc::getrandom(buf.as_mut_ptr().cast(), buf.len(), 0)
+        let ret = libc::getrandom(buf.as_mut_ptr().cast(), buf.len(), 0);
+
+        #[cfg(any(target_os = "android", target_os = "linux"))]
+        #[allow(unused_unsafe)] // TODO(MSRV 1.65): Remove this.
+        unsafe {
+            super::sanitizer::unpoison_linux_getrandom_result(buf, ret);
+        }
+
+        ret
     })
 }

src/backends/linux_android_with_fallback.rs

@@ -1,5 +1,5 @@
 //! Implementation for Linux / Android with `/dev/urandom` fallback
-use super::use_file;
+use super::{sanitizer, use_file};
 use crate::Error;
 use core::{
     ffi::c_void,
@@ -95,7 +95,9 @@ pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
         // note: `transmute` is currently the only way to convert a pointer into a function reference
         let getrandom_fn = unsafe { transmute::<NonNull<c_void>, GetRandomFn>(fptr) };
         util_libc::sys_fill_exact(dest, |buf| unsafe {
-            getrandom_fn(buf.as_mut_ptr().cast(), buf.len(), 0)
+            let ret = getrandom_fn(buf.as_mut_ptr().cast(), buf.len(), 0);
+            sanitizer::unpoison_linux_getrandom_result(buf, ret);
+            ret
         })
     }
 }

src/backends/linux_raw.rs

@@ -1,7 +1,7 @@
 //! Implementation for Linux / Android using `asm!`-based syscalls.
-use crate::{Error, MaybeUninit};
-
+use super::sanitizer;
 pub use crate::util::{inner_u32, inner_u64};
+use crate::{Error, MaybeUninit};
 
 #[cfg(not(any(target_os = "android", target_os = "linux")))]
 compile_error!("`linux_raw` backend can be enabled only for Linux/Android targets!");
@@ -118,6 +118,7 @@ pub fn fill_inner(mut dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
 
     loop {
         let ret = unsafe { getrandom_syscall(dest.as_mut_ptr().cast(), dest.len(), 0) };
+        unsafe { sanitizer::unpoison_linux_getrandom_result(dest, ret) };
         match usize::try_from(ret) {
             Ok(0) => return Err(Error::UNEXPECTED),
             Ok(len) => {

src/backends/sanitizer.rs

@@ -0,0 +1,59 @@
+use core::mem::MaybeUninit;
+
+/// Unpoisons `buf` if MSAN support is enabled.
+///
+/// Most backends do not need to unpoison their output. Rust language- and
+/// library- provided functionality unpoisons automatically. Similarly, libc
+/// either natively supports MSAN and/or MSAN hooks libc-provided functions
+/// to unpoison outputs on success. Only when all of these things are
+/// bypassed do we need to do it ourselves.
+///
+/// The call to unpoison should be done as close to the write as possible.
+/// For example, if the backend partially fills the output buffer in chunks,
+/// each chunk should be unpoisoned individually. This way, the correctness of
+/// the chunking logic can be validated (in part) using MSAN.
+pub unsafe fn unpoison(buf: &mut [MaybeUninit<u8>]) {
+    cfg_if! {
+        if #[cfg(getrandom_msan)] {
+            extern "C" {
+                fn __msan_unpoison(a: *mut core::ffi::c_void, size: usize);
+            }
+            let a = buf.as_mut_ptr().cast();
+            let size = buf.len();
+            #[allow(unused_unsafe)] // TODO(MSRV 1.65): Remove this.
+            unsafe {
+                __msan_unpoison(a, size);
+            }
+        } else {
+            let _ = buf;
+        }
+    }
+}
+
+/// Interprets the result of the `getrandom` syscall of Linux, unpoisoning any
+/// written part of `buf`.
+///
+/// `buf` must be the output buffer that was originally passed to the `getrandom`
+/// syscall.
+///
+/// `ret` must be the result returned by `getrandom`. If `ret` is negative or
+/// larger than the length of `buf` then nothing is done.
+///
+/// Memory Sanitizer only intercepts `getrandom` on this condition (from its
+/// source code):
+/// ```c
+/// #define SANITIZER_INTERCEPT_GETRANDOM \
+///   ((SI_LINUX && __GLIBC_PREREQ(2, 25)) || SI_FREEBSD || SI_SOLARIS)
+/// ```
+/// So, effectively, we have to assume that it is never intercepted on Linux.
+#[cfg(any(target_os = "android", target_os = "linux"))]
+pub unsafe fn unpoison_linux_getrandom_result(buf: &mut [MaybeUninit<u8>], ret: isize) {
+    if let Ok(bytes_written) = usize::try_from(ret) {
+        if let Some(written) = buf.get_mut(..bytes_written) {
+            #[allow(unused_unsafe)] // TODO(MSRV 1.65): Remove this.
+            unsafe {
+                unpoison(written)
+            }
+        }
+    }
+}

src/lib.rs

@@ -106,12 +106,7 @@ pub fn fill_uninit(dest: &mut [MaybeUninit<u8>]) -> Result<&mut [u8], Error> {
 
     // SAFETY: `dest` has been fully initialized by `imp::fill_inner`
     // since it returned `Ok`.
-    Ok(unsafe {
-        #[cfg(getrandom_msan)]
-        __msan_unpoison(dest.as_mut_ptr().cast(), dest.len());
-
-        util::slice_assume_init_mut(dest)
-    })
+    Ok(unsafe { util::slice_assume_init_mut(dest) })
 }
 
 /// Get random `u32` from the system's preferred random number source.