Add TryFromBytes::try_from_mut (google#892)

joshlf · dorryspears · commit 824d7dc14140 · 2024-02-19T21:22:27.000-05:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -1269,11 +1269,43 @@ pub unsafe trait TryFromBytes {
         // This call may panic. If that happens, it doesn't cause any soundness
         // issues, as we have not generated any invalid state which we need to
         // fix before returning.
-        let candidate = candidate.check_valid();
+        let candidate = candidate.try_into_valid();
 
         candidate.map(MaybeAligned::as_ref)
     }
 
+    /// Attempts to interpret a mutable byte slice as a `Self`.
+    ///
+    /// `try_from_mut` validates that `bytes` contains a valid `Self`, and that
+    /// it satisfies `Self`'s alignment requirement. If it does, then `bytes` is
+    /// reinterpreted as a `Self`.
+    ///
+    /// Note that Rust's bit validity rules are still being decided. As such,
+    /// there exist types whose bit validity is ambiguous. See
+    /// [here][TryFromBytes#what-is-a-valid-instance] for a discussion of how
+    /// these cases are handled.
+    // TODO(#251): Require `Self: NoCell` and allow `TryFromBytes` types to
+    // contain `UnsafeCell`s.
+    #[inline]
+    #[doc(hidden)] // TODO(#5): Finalize name before remove this attribute.
+    fn try_from_mut(bytes: &mut [u8]) -> Option<&mut Self>
+    where
+        Self: KnownLayout + NoCell, // TODO(#251): Remove the `NoCell` bound.
+    {
+        let candidate = Ptr::from_mut(bytes).try_cast_into_no_leftover::<Self>()?;
+
+        // SAFETY: `candidate` has no uninitialized sub-ranges because it
+        // derived from `bytes: &mut [u8]`.
+        let candidate = unsafe { candidate.assume_validity::<invariant::Initialized>() };
+
+        // This call may panic. If that happens, it doesn't cause any soundness
+        // issues, as we have not generated any invalid state which we need to
+        // fix before returning.
+        let candidate = candidate.try_into_valid();
+
+        candidate.map(Ptr::as_mut)
+    }
+
     /// Attempts to read a `Self` from a byte slice.
     ///
     /// `try_read_from` validates that `bytes` contains a valid `Self`. If it
@@ -7985,7 +8017,7 @@ mod tests {
         // `impl_try_from_bytes_testable!`).
         trait TryFromBytesTestable {
             fn with_passing_test_cases<F: Fn(&Self)>(f: F);
-            fn with_failing_test_cases<F: Fn(&[u8])>(f: F);
+            fn with_failing_test_cases<F: Fn(&mut [u8])>(f: F);
         }
 
         impl<T: FromBytes> TryFromBytesTestable for T {
@@ -8005,7 +8037,7 @@ mod tests {
                 f(&ffs);
             }
 
-            fn with_failing_test_cases<F: Fn(&[u8])>(_f: F) {}
+            fn with_failing_test_cases<F: Fn(&mut [u8])>(_f: F) {}
         }
 
         // Implements `TryFromBytesTestable`.
@@ -8037,13 +8069,13 @@ mod tests {
                     )*
                 }
 
-                fn with_failing_test_cases<F: Fn(&[u8])>(_f: F) {
+                fn with_failing_test_cases<F: Fn(&mut [u8])>(_f: F) {
                     $($(
                         // `unused_qualifications` is spuriously triggered on
                         // `Option::<Self>::None`.
                         #[allow(unused_qualifications)]
-                        let case = $failure_case.as_bytes();
-                        _f(case.as_bytes());
+                        let mut case = $failure_case;//.as_mut_bytes();
+                        _f(case.as_mut_bytes());
                     )*)?
                 }
             };
@@ -8135,7 +8167,9 @@ mod tests {
                 #[allow(clippy::as_conversions)]
                 <$ty as TryFromBytesTestable>::with_failing_test_cases(|c| {
                     let res = <$ty as TryFromBytes>::try_from_ref(c);
-                    assert!(res.is_none(), "{}::is_bit_valid({:?}): got true, expected false", stringify!($ty), c);
+                    assert!(res.is_none(), "{}::try_from_ref({:?}): got Some, expected None", stringify!($ty), c);
+                    let res = <$ty as TryFromBytes>::try_from_mut(c);
+                    assert!(res.is_none(), "{}::try_from_mut({:?}): got Some, expected None", stringify!($ty), c);
                 });
 
                 #[allow(dead_code)]
diff --git a/src/pointer/mod.rs b/src/pointer/mod.rs
@@ -12,43 +12,19 @@ mod ptr;
 
 pub use ptr::{invariant, Ptr};
 
-use crate::{TryFromBytes, Unaligned};
+use crate::Unaligned;
 
 /// A shorthand for a maybe-valid, maybe-aligned reference. Used as the argument
 /// to [`TryFromBytes::is_bit_valid`].
+///
+/// [`TryFromBytes::is_bit_valid`]: crate::TryFromBytes::is_bit_valid
 pub type Maybe<'a, T, Alignment = invariant::AnyAlignment> =
     Ptr<'a, T, (invariant::Shared, Alignment, invariant::Initialized)>;
 
-// These methods are defined on the type alias, `Maybe`, so as to bring them to
-// the forefront of the rendered rustdoc.
-impl<'a, T, Alignment> Maybe<'a, T, Alignment>
-where
-    T: 'a + ?Sized + TryFromBytes,
-    Alignment: invariant::Alignment,
-{
-    /// Checks that `Ptr`'s referent is validly initialized for `T`.
-    ///
-    /// # Panics
-    ///
-    /// This method will panic if
-    /// [`T::is_bit_valid`][TryFromBytes::is_bit_valid] panics.
-    #[inline]
-    pub(crate) fn check_valid(self) -> Option<MaybeAligned<'a, T, Alignment>> {
-        // This call may panic. If that happens, it doesn't cause any soundness
-        // issues, as we have not generated any invalid state which we need to
-        // fix before returning.
-        if T::is_bit_valid(self.forget_aligned()) {
-            // SAFETY: If `T::is_bit_valid`, code may assume that `self`
-            // contains a bit-valid instance of `Self`.
-            Some(unsafe { self.assume_validity::<invariant::Valid>() })
-        } else {
-            None
-        }
-    }
-}
-
 /// A semi-user-facing wrapper type representing a maybe-aligned reference, for
 /// use in [`TryFromBytes::is_bit_valid`].
+///
+/// [`TryFromBytes::is_bit_valid`]: crate::TryFromBytes::is_bit_valid
 pub type MaybeAligned<'a, T, Alignment = invariant::AnyAlignment> =
     Ptr<'a, T, (invariant::Shared, Alignment, invariant::Valid)>;
 
diff --git a/src/pointer/ptr.rs b/src/pointer/ptr.rs
@@ -419,6 +419,38 @@ mod _conversions {
         }
     }
 
+    /// `&'a mut T` → `Ptr<'a, T>`
+    impl<'a, T> Ptr<'a, T, (invariant::Exclusive, invariant::Aligned, invariant::Valid)>
+    where
+        T: 'a + ?Sized,
+    {
+        /// Constructs a `Ptr` from an exclusive reference.
+        #[inline]
+        pub(crate) fn from_mut(ptr: &'a mut T) -> Self {
+            let ptr = core::ptr::NonNull::from(ptr);
+            // SAFETY:
+            // 0. `ptr`, by invariant on `&'a mut T`, is derived from some valid
+            //    Rust allocation, `A`.
+            // 1. `ptr`, by invariant on `&'a mut T`, has valid provenance for
+            //    `A`.
+            // 2. `ptr`, by invariant on `&'a mut T`, addresses a byte range
+            //    which is entirely contained in `A`.
+            // 3. `ptr`, by invariant on `&'a mut T`, addresses a byte range
+            //    whose length fits in an `isize`.
+            // 4. `ptr`, by invariant on `&'a mut T`, addresses a byte range
+            //     which does not wrap around the address space.
+            // 5. `A`, by invariant on `&'a mut T`, is guaranteed to live for at
+            //    least `'a`.
+            // 6. `ptr`, by invariant on `&'a mut T`, conforms to the aliasing
+            //    invariant of `invariant::Exclusive`.
+            // 7. `ptr`, by invariant on `&'a mut T`, conforms to the alignment
+            //    invariant of `invariant::Aligned`.
+            // 8. `ptr`, by invariant on `&'a mut T`, conforms to the validity
+            //    invariant of `invariant::Valid`.
+            unsafe { Self::new(ptr) }
+        }
+    }
+
     /// `Ptr<'a, T>` → `&'a T`
     impl<'a, T> Ptr<'a, T, (invariant::Shared, invariant::Aligned, invariant::Valid)>
     where
@@ -460,6 +492,104 @@ mod _conversions {
         }
     }
 
+    impl<'a, T, I> Ptr<'a, T, I>
+    where
+        T: 'a + ?Sized,
+        I: Invariants,
+        I::Aliasing: invariant::at_least::Shared,
+    {
+        /// Reborrows this `Ptr`, producing another `Ptr`.
+        ///
+        /// Since `self` is borrowed immutably, this prevents any mutable
+        /// methods from being called on `self` as long as the returned `Ptr`
+        /// exists.
+        #[allow(clippy::needless_lifetimes)] // Allows us to name the lifetime in the safety comment below.
+        pub(crate) fn reborrow<'b>(&'b mut self) -> Ptr<'b, T, I>
+        where
+            'a: 'b,
+        {
+            // SAFETY: The following all hold by invariant on `self`, and thus
+            // hold of `ptr = self.as_non_null()`:
+            // 0. `ptr` is derived from some valid Rust allocation, `A`.
+            // 1. `ptr` has valid provenance for `A`.
+            // 2. `ptr` addresses a byte range which is entirely contained in
+            //    `A`.
+            // 3. `ptr` addresses a byte range whose length fits in an `isize`.
+            // 4. `ptr` addresses a byte range which does not wrap around the
+            //    address space.
+            // 5. `A` is guaranteed to live for at least `'a`.
+            // 6. SEE BELOW.
+            // 7. `ptr` conforms to the alignment invariant of
+            //   [`I::Alignment`](invariant::Alignment).
+            // 8. `ptr` conforms to the validity invariant of
+            //   [`I::Validity`](invariant::Validity).
+            // 9. During the lifetime 'a, no code will load or store this memory
+            //    region treating `UnsafeCell`s as existing at different ranges
+            //    than they exist in `T`.
+            // 10. During the lifetime 'a, no reference will exist to this
+            //     memory which treats `UnsafeCell`s as existing at different
+            //     ranges than they exist in `T`.
+            //
+            // For aliasing (6 above), since `I::Aliasing:
+            // invariant::at_least::Shared`, there are two cases for
+            // `I::Aliasing`:
+            // - For `invariant::Shared`: `'a` outlives `'b`, and so the
+            //   returned `Ptr` does not permit accessing the referent any
+            //   longer than is possible via `self`. For shared aliasing, it is
+            //   sound for multiple `Ptr`s to exist simultaneously which
+            //   reference the same memory, so creating a new one is not
+            //   problematic.
+            // - For `invariant::Exclusive`: Since `self` is `&'b mut` and we
+            //   return a `Ptr` with lifetime `'b`, `self` is inaccessible to
+            //   the caller for the lifetime `'b` - in other words, `self` is
+            //   inaccessible to the caller as long as the returned `Ptr`
+            //   exists. Since `self` is an exclusive `Ptr`, no other live
+            //   references or `Ptr`s may exist which refer to the same memory
+            //   while `self` is live. Thus, as long as the returned `Ptr`
+            //   exists, no other references or `Ptr`s which refer to the same
+            //   memory may be live.
+            unsafe { Ptr::new(self.as_non_null()) }
+        }
+    }
+
+    /// `Ptr<'a, T>` → `&'a mut T`
+    impl<'a, T> Ptr<'a, T, (invariant::Exclusive, invariant::Aligned, invariant::Valid)>
+    where
+        T: 'a + ?Sized,
+    {
+        /// Converts the `Ptr` to a mutable reference.
+        #[allow(clippy::wrong_self_convention)]
+        pub(crate) fn as_mut(self) -> &'a mut T {
+            let mut raw = self.as_non_null();
+            // SAFETY: This invocation of `NonNull::as_mut` satisfies its
+            // documented safety preconditions:
+            //
+            // 1. The pointer is properly aligned. This is ensured by-contract
+            //    on `Ptr`, because the `ALIGNMENT_INVARIANT` is `Aligned`.
+            //
+            // 2. It must be “dereferenceable” in the sense defined in the
+            //    module documentation; i.e.:
+            //
+            //    > The memory range of the given size starting at the pointer
+            //    > must all be within the bounds of a single allocated object.
+            //    > [2]
+            //
+            //   This is ensured by contract on all `Ptr`s.
+            //
+            // 3. The pointer must point to an initialized instance of `T`. This
+            //    is ensured by-contract on `Ptr`, because the
+            //    `VALIDITY_INVARIANT` is `Valid`.
+            //
+            // 4. You must enforce Rust’s aliasing rules. This is ensured by
+            //    contract on `Ptr`, because the `ALIASING_INVARIANT` is
+            //    `Exclusive`.
+            //
+            // [1]: https://doc.rust-lang.org/std/ptr/struct.NonNull.html#method.as_mut
+            // [2]: https://doc.rust-lang.org/std/ptr/index.html#safety
+            unsafe { raw.as_mut() }
+        }
+    }
+
     /// `&'a MaybeUninit<T>` → `Ptr<'a, T>`
     impl<'a, T> Ptr<'a, T, (invariant::Shared, invariant::Aligned, invariant::AnyValidity)>
     where
@@ -506,6 +636,7 @@ mod _conversions {
 /// State transitions between invariants.
 mod _transitions {
     use super::*;
+    use crate::TryFromBytes;
 
     impl<'a, T, I> Ptr<'a, T, I>
     where
@@ -592,6 +723,46 @@ mod _transitions {
             unsafe { self.assume_validity::<invariant::Valid>() }
         }
 
+        /// Checks that `Ptr`'s referent is validly initialized for `T`,
+        /// returning a `Ptr` with `invariant::Valid` on success.
+        ///
+        /// # Panics
+        ///
+        /// This method will panic if
+        /// [`T::is_bit_valid`][TryFromBytes::is_bit_valid] panics.
+        #[inline]
+        pub(crate) fn try_into_valid(
+            mut self,
+        ) -> Option<Ptr<'a, T, (I::Aliasing, I::Alignment, invariant::Valid)>>
+        where
+            T: TryFromBytes,
+            I::Aliasing: invariant::at_least::Shared,
+            I: Invariants<Validity = invariant::Initialized>,
+        {
+            // This call may panic. If that happens, it doesn't cause any soundness
+            // issues, as we have not generated any invalid state which we need to
+            // fix before returning.
+            if T::is_bit_valid(self.reborrow().forget_exclusive().forget_aligned()) {
+                // SAFETY: If `T::is_bit_valid`, code may assume that `self`
+                // contains a bit-valid instance of `Self`.
+                Some(unsafe { self.assume_validity::<invariant::Valid>() })
+            } else {
+                None
+            }
+        }
+
+        /// Forgets that `Ptr`'s referent exclusively references `T`,
+        /// downgrading to a shared reference.
+        #[doc(hidden)]
+        #[inline]
+        pub fn forget_exclusive(self) -> Ptr<'a, T, (invariant::Shared, I::Alignment, I::Validity)>
+        where
+            I::Aliasing: invariant::at_least::Shared,
+        {
+            // SAFETY: `I::Aliasing` is at least as restrictive as `Shared`.
+            unsafe { Ptr::from_ptr(self) }
+        }
+
         /// Forgets that `Ptr`'s referent is validly-aligned for `T`.
         #[doc(hidden)]
         #[inline]
