From b329c85e6c399269102aa6eb389ae6b320b2f73a Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 5 Apr 2024 13:09:23 -0700
Subject: [PATCH] Implement check for `write_bytes` (#3108)

In the previous PR #3085, we did not support checks for `write_bytes`
which is added in this PR.

I am waiting for #3092 to add expected tests.
---
 .../src/kani_middle/transform/check_values.rs | 21 ++++++++++++++++++-
 tests/kani/ValidValues/maybe_uninit.rs        | 21 +++++++++++++++++++
 tests/kani/ValidValues/write_bytes.rs         | 21 +++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 tests/kani/ValidValues/maybe_uninit.rs
 create mode 100644 tests/kani/ValidValues/write_bytes.rs

diff --git a/kani-compiler/src/kani_middle/transform/check_values.rs b/kani-compiler/src/kani_middle/transform/check_values.rs
index 4888e032af4f..a620dfabb72e 100644
--- a/kani-compiler/src/kani_middle/transform/check_values.rs
+++ b/kani-compiler/src/kani_middle/transform/check_values.rs
@@ -440,7 +440,13 @@ impl<'a> MirVisitor for CheckValueVisitor<'a> {
                                 // The write bytes intrinsic may trigger UB in safe code.
                                 // pub unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize)
                                 // <https://doc.rust-lang.org/stable/core/intrinsics/fn.write_bytes.html>
-                                // We don't support this operation yet.
+                                // This is an over-approximation since writing an invalid value is
+                                // not UB, only reading it will be.
+                                assert_eq!(
+                                    args.len(),
+                                    3,
+                                    "Unexpected number of arguments for `write_bytes`"
+                                );
                                 let TyKind::RigidTy(RigidTy::RawPtr(target_ty, Mutability::Mut)) =
                                     args[0].ty(self.locals).unwrap().kind()
                                 else {
@@ -449,6 +455,19 @@ impl<'a> MirVisitor for CheckValueVisitor<'a> {
                                 let validity = ty_validity_per_offset(&self.machine, target_ty, 0);
                                 match validity {
                                     Ok(ranges) if ranges.is_empty() => {}
+                                    Ok(ranges) => {
+                                        let sz = Const::try_from_uint(
+                                            target_ty.layout().unwrap().shape().size.bytes()
+                                                as u128,
+                                            UintTy::Usize,
+                                        )
+                                        .unwrap();
+                                        self.push_target(SourceOp::BytesValidity {
+                                            target_ty,
+                                            rvalue: Rvalue::Repeat(args[1].clone(), sz),
+                                            ranges,
+                                        })
+                                    }
                                     _ => self.push_target(SourceOp::UnsupportedCheck {
                                         check: "write_bytes".to_string(),
                                         ty: target_ty,
diff --git a/tests/kani/ValidValues/maybe_uninit.rs b/tests/kani/ValidValues/maybe_uninit.rs
new file mode 100644
index 000000000000..620ad8ef5ba3
--- /dev/null
+++ b/tests/kani/ValidValues/maybe_uninit.rs
@@ -0,0 +1,21 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z valid-value-checks
+//! Check that Kani can identify UB when converting from a maybe uninit().
+
+use std::mem::MaybeUninit;
+use std::num::NonZeroI64;
+
+#[kani::proof]
+pub fn check_valid_zeroed() {
+    let maybe = MaybeUninit::zeroed();
+    let val: u128 = unsafe { maybe.assume_init() };
+    assert_eq!(val, 0);
+}
+
+#[kani::proof]
+#[kani::should_panic]
+pub fn check_invalid_zeroed() {
+    let maybe = MaybeUninit::zeroed();
+    let _val: NonZeroI64 = unsafe { maybe.assume_init() };
+}
diff --git a/tests/kani/ValidValues/write_bytes.rs b/tests/kani/ValidValues/write_bytes.rs
new file mode 100644
index 000000000000..e4f73b1f3479
--- /dev/null
+++ b/tests/kani/ValidValues/write_bytes.rs
@@ -0,0 +1,21 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z valid-value-checks
+//! Check that Kani can identify UB when using write_bytes for initializing a variable.
+#![feature(core_intrinsics)]
+
+#[kani::proof]
+#[kani::should_panic]
+pub fn check_invalid_write() {
+    let mut val = 'a';
+    let ptr = &mut val as *mut char;
+    // Should fail given that we wrote invalid value to array of char.
+    unsafe { std::intrinsics::write_bytes(ptr, kani::any(), 1) };
+}
+
+#[kani::proof]
+pub fn check_valid_write() {
+    let mut val = 10u128;
+    let ptr = &mut val as *mut _;
+    unsafe { std::intrinsics::write_bytes(ptr, kani::any(), 1) };
+}