Initial work on runtime type validity checks

Author: 5225225

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -106,7 +106,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             | sym::needs_drop
             | sym::type_id
             | sym::type_name
-            | sym::variant_count => {
+            | sym::variant_count
+            | sym::validity_invariants_of => {
                 let value = bx
                     .tcx()
                     .const_eval_instance(ty::ParamEnv::reveal_all(), instance, None)

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -24,6 +24,7 @@ use super::{
 
 mod caller_location;
 mod type_name;
+mod validity_invariants_of;
 
 fn numeric_intrinsic<Tag>(name: Symbol, bits: u128, kind: Primitive) -> Scalar<Tag> {
     let size = match kind {
@@ -103,6 +104,11 @@ pub(crate) fn eval_nullary_intrinsic<'tcx>(
             | ty::Tuple(_)
             | ty::Error(_) => ConstValue::from_machine_usize(0u64, &tcx),
         },
+        sym::validity_invariants_of => {
+            ensure_monomorphic_enough(tcx, tp_ty)?;
+            let alloc = validity_invariants_of::alloc_validity_invariants_of(tcx, tp_ty);
+            ConstValue::Slice { data: alloc, start: 0, end: alloc.inner().len() }
+        }
         other => bug!("`{}` is not a zero arg intrinsic", other),
     })
 }
@@ -162,13 +168,15 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | sym::needs_drop
             | sym::type_id
             | sym::type_name
+            | sym::validity_invariants_of
             | sym::variant_count => {
                 let gid = GlobalId { instance, promoted: None };
                 let ty = match intrinsic_name {
                     sym::pref_align_of | sym::variant_count => self.tcx.types.usize,
                     sym::needs_drop => self.tcx.types.bool,
                     sym::type_id => self.tcx.types.u64,
                     sym::type_name => self.tcx.mk_static_str(),
+                    sym::validity_invariants_of => self.tcx.mk_static_bytes(),
                     _ => bug!(),
                 };
                 let val =

compiler/rustc_const_eval/src/interpret/intrinsics/validity_invariants_of.rs

@@ -0,0 +1,119 @@
+use rustc_middle::mir::interpret::{Allocation, ConstAllocation};
+use rustc_middle::mir::Mutability;
+use rustc_middle::ty::layout::LayoutCx;
+use rustc_middle::ty::{ParamEnv, ParamEnvAnd};
+use rustc_middle::ty::{Ty, TyCtxt};
+use rustc_target::abi::{
+    Abi, Align, Endian, FieldsShape, HasDataLayout, Scalar, Size, WrappingRange,
+};
+
+#[derive(Debug, Clone, Copy)]
+struct Invariant {
+    offset: Size,
+    size: Size,
+    start: u128,
+    end: u128,
+}
+
+fn add_invariants<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, invs: &mut Vec<Invariant>, offset: Size) {
+    let x = tcx.layout_of(ParamEnvAnd { param_env: ParamEnv::reveal_all(), value: ty });
+
+    if let Ok(layout) = x {
+        if let Abi::Scalar(Scalar::Initialized { value, valid_range }) = layout.layout.abi() {
+            let size = value.size(&tcx);
+            let WrappingRange { start, end } = valid_range;
+            invs.push(Invariant { offset, size, start, end })
+        }
+
+        let param_env = ParamEnv::reveal_all();
+        let unwrap = LayoutCx { tcx, param_env };
+
+        match layout.layout.fields() {
+            FieldsShape::Primitive => {}
+            FieldsShape::Union(_) => {}
+            FieldsShape::Array { stride, count } => {
+                // TODO: should we just bail if we're making a Too Large type?
+                // (Like [bool; 1_000_000])
+                for idx in 0..*count {
+                    let off = offset + *stride * idx;
+                    let f = layout.field(&unwrap, idx as usize);
+                    add_invariants(tcx, f.ty, invs, off);
+                }
+            }
+            FieldsShape::Arbitrary { offsets, .. } => {
+                for (idx, &field_offset) in offsets.iter().enumerate() {
+                    let f = layout.field(&unwrap, idx);
+                    if f.ty == ty {
+                        // Some types contain themselves as fields, such as
+                        // &mut [T]
+                        // Easy solution is to just not recurse then.
+                    } else {
+                        add_invariants(tcx, f.ty, invs, offset + field_offset);
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn extend_encoded_int(to: &mut Vec<u8>, endian: Endian, ptr_size: PointerSize, value: Size) {
+    match (endian, ptr_size) {
+        (Endian::Little, PointerSize::Bits16) => to.extend((value.bytes() as u16).to_le_bytes()),
+        (Endian::Little, PointerSize::Bits32) => to.extend((value.bytes() as u32).to_le_bytes()),
+        (Endian::Little, PointerSize::Bits64) => to.extend((value.bytes()).to_le_bytes()),
+        (Endian::Big, PointerSize::Bits16) => to.extend((value.bytes() as u16).to_be_bytes()),
+        (Endian::Big, PointerSize::Bits32) => to.extend((value.bytes() as u32).to_be_bytes()),
+        (Endian::Big, PointerSize::Bits64) => to.extend((value.bytes()).to_be_bytes()),
+    }
+}
+
+#[derive(Clone, Copy)]
+enum PointerSize {
+    Bits16,
+    Bits32,
+    Bits64,
+}
+
+/// Directly returns a `ConstAllocation` containing a list of validity invariants of the given type.
+pub(crate) fn alloc_validity_invariants_of<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    ty: Ty<'tcx>,
+) -> ConstAllocation<'tcx> {
+    let mut invs: Vec<Invariant> = Vec::new();
+
+    let layout = tcx.data_layout();
+
+    let ptr_size = match layout.pointer_size.bits() {
+        16 => PointerSize::Bits16,
+        32 => PointerSize::Bits32,
+        64 => PointerSize::Bits64,
+        _ => {
+            // Not sure if this can happen, but just return an empty slice?
+            let alloc =
+                Allocation::from_bytes(Vec::new(), Align::from_bytes(8).unwrap(), Mutability::Not);
+            return tcx.intern_const_alloc(alloc);
+        }
+    };
+
+    add_invariants(tcx, ty, &mut invs, Size::ZERO);
+
+    let encode_range = match layout.endian {
+        Endian::Little => |r: u128| r.to_le_bytes(),
+        Endian::Big => |r: u128| r.to_be_bytes(),
+    };
+
+    let mut encoded = Vec::new();
+
+    // TODO: this needs to match the layout of `Invariant` in core/src/intrinsics.rs
+    // how do we ensure that?
+    for inv in invs {
+        extend_encoded_int(&mut encoded, layout.endian, ptr_size, inv.offset);
+        extend_encoded_int(&mut encoded, layout.endian, ptr_size, inv.size);
+        encoded.extend(encode_range(inv.start));
+        encoded.extend(encode_range(inv.end));
+    }
+
+    // TODO: The alignment here should be calculated from the struct definition, I guess?
+    let alloc = Allocation::from_bytes(encoded, Align::from_bytes(8).unwrap(), Mutability::Not);
+    tcx.intern_const_alloc(alloc)
+}

compiler/rustc_middle/src/ty/context.rs

@@ -2339,6 +2339,11 @@ impl<'tcx> TyCtxt<'tcx> {
         self.mk_imm_ref(self.lifetimes.re_static, self.types.str_)
     }
 
+    #[inline]
+    pub fn mk_static_bytes(self) -> Ty<'tcx> {
+        self.mk_imm_ref(self.lifetimes.re_static, self.mk_slice(self.types.u8))
+    }
+
     #[inline]
     pub fn mk_adt(self, def: AdtDef<'tcx>, substs: SubstsRef<'tcx>) -> Ty<'tcx> {
         // Take a copy of substs so that we own the vectors inside.

compiler/rustc_span/src/symbol.rs

@@ -1524,6 +1524,7 @@ symbols! {
         va_list,
         va_start,
         val,
+        validity_invariants_of,
         values,
         var,
         variant_count,

compiler/rustc_typeck/src/check/intrinsic.rs

@@ -102,7 +102,8 @@ pub fn intrinsic_operation_unsafety(intrinsic: Symbol) -> hir::Unsafety {
         | sym::type_name
         | sym::forget
         | sym::black_box
-        | sym::variant_count => hir::Unsafety::Normal,
+        | sym::variant_count
+        | sym::validity_invariants_of => hir::Unsafety::Normal,
         _ => hir::Unsafety::Unsafe,
     }
 }
@@ -191,6 +192,7 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
             sym::needs_drop => (1, Vec::new(), tcx.types.bool),
 
             sym::type_name => (1, Vec::new(), tcx.mk_static_str()),
+            sym::validity_invariants_of => (1, Vec::new(), tcx.mk_static_bytes()),
             sym::type_id => (1, Vec::new(), tcx.types.u64),
             sym::offset | sym::arith_offset => (
                 1,

library/core/src/intrinsics.rs

@@ -2138,6 +2138,30 @@ pub const unsafe fn copy_nonoverlapping<T>(src: *const T, dst: *mut T, count: us
     }
 }
 
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct Invariant {
+    offset: u64,
+    size: u64,
+    start: u128,
+    end: u128,
+}
+
+#[cfg(not(bootstrap))]
+/// Returns a list of all validity invariants of the type.
+pub const fn validity_invariants_of<T>() -> &'static [Invariant] {
+    extern "rust-intrinsic" {
+        #[rustc_const_stable(feature = "validity_invariants_of", since = "1.40.0")]
+        pub fn validity_invariants_of<T>() -> &'static [u8];
+    }
+
+    let invariants: &'static [u8] = validity_invariants_of::<T>();
+    let sz = invariants.len() / core::mem::size_of::<Invariant>();
+
+    // SAFETY: we know this is valid.
+    unsafe { core::slice::from_raw_parts(invariants.as_ptr().cast(), sz) }
+}
+
 /// Copies `count * size_of::<T>()` bytes from `src` to `dst`. The source
 /// and destination may overlap.
 ///
@@ -2394,3 +2418,59 @@ where
 {
     called_in_const.call_once(arg)
 }
+
+#[cfg(bootstrap)]
+pub(crate) const unsafe fn assert_validity_of<T>(_: *const T) -> bool {
+    true
+}
+
+#[cfg(not(bootstrap))]
+/// Asserts that the value at `value` is valid at type T.
+/// Best effort, and is UB if the value is invalid.
+pub(crate) unsafe fn assert_validity_of<T>(value: *const T) -> bool {
+    #[repr(packed)]
+    struct Unaligned<T>(T);
+
+    // SAFETY:
+    unsafe {
+        let invariants = validity_invariants_of::<T>();
+        for invariant in invariants {
+            let off = invariant.offset as usize;
+            let start = invariant.start;
+            let end = invariant.end;
+
+            let (value, max): (u128, u128) = match invariant.size {
+                1 => ((*(value.cast::<u8>().add(off))).into(), u8::MAX.into()),
+                2 => (
+                    (*value.cast::<u8>().add(off).cast::<Unaligned<u16>>()).0.into(),
+                    u16::MAX.into(),
+                ),
+                4 => (
+                    (*value.cast::<u8>().add(off).cast::<Unaligned<u32>>()).0.into(),
+                    u32::MAX.into(),
+                ),
+                8 => (
+                    (*value.cast::<u8>().add(off).cast::<Unaligned<u64>>()).0.into(),
+                    u64::MAX.into(),
+                ),
+                16 => (
+                    (*value.cast::<u8>().add(off).cast::<Unaligned<u128>>()).0.into(),
+                    u128::MAX.into(),
+                ),
+                s => panic!("unexpected size {s}"),
+            };
+
+            if start > end {
+                if !((start..=max).contains(&value) || (0..=end).contains(&value)) {
+                    return false;
+                }
+            } else {
+                if !(start..=end).contains(&value) {
+                    return false;
+                }
+            }
+        }
+
+        true
+    }
+}

library/core/src/mem/maybe_uninit.rs

@@ -631,6 +631,11 @@ impl<T> MaybeUninit<T> {
         // This also means that `self` must be a `value` variant.
         unsafe {
             intrinsics::assert_inhabited::<T>();
+
+            intrinsics::assert_unsafe_precondition!(intrinsics::assert_validity_of::<T>(
+                &self as *const MaybeUninit<T> as *const T
+            ));
+
             ManuallyDrop::into_inner(self.value)
         }
     }
@@ -795,6 +800,11 @@ impl<T> MaybeUninit<T> {
         // This also means that `self` must be a `value` variant.
         unsafe {
             intrinsics::assert_inhabited::<T>();
+
+            intrinsics::assert_unsafe_precondition!(intrinsics::assert_validity_of::<T>(
+                self as *const MaybeUninit<T> as *const T
+            ));
+
             &*self.as_ptr()
         }
     }
@@ -912,6 +922,11 @@ impl<T> MaybeUninit<T> {
         // This also means that `self` must be a `value` variant.
         unsafe {
             intrinsics::assert_inhabited::<T>();
+
+            intrinsics::assert_unsafe_precondition!(intrinsics::assert_validity_of::<T>(
+                self as *const MaybeUninit<T> as *const T
+            ));
+
             &mut *self.as_mut_ptr()
         }
     }