cmse: clear variant-dependent padding in enums

compiler/rustc_abi/src/layout/ty.rs

@@ -284,14 +284,35 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
         found
     }
 
+    /// Whether this type/layout has any padding that is dependent on a variant, i.e. has bytes that
+    /// are padding for some, but not all, valid values of this type.
+    pub fn has_variant_dependent_padding<C>(&self, cx: &C) -> bool
+    where
+        Ty: TyAbiInterface<'a, C> + Copy,
+    {
+        match self.variants {
+            Variants::Multiple { .. } => true,
+            Variants::Empty => false,
+            Variants::Single { .. } => match &self.fields {
+                FieldsShape::Primitive | FieldsShape::Union(_) => false,
+                FieldsShape::Array { count, .. } => {
+                    *count > 0 && self.field(cx, 0).has_variant_dependent_padding(cx)
+                }
+                FieldsShape::Arbitrary { offsets, .. } => {
+                    (0..offsets.len()).any(|i| self.field(cx, i).has_variant_dependent_padding(cx))
+                }
+            },
+        }
+    }
+
     /// The ranges of bytes that are always ignored by the representation relation of this type.
     ///
     /// In other words, for any sequence of bytes, if we reset the these padding bytes to uninit,
     /// then these two sequences of bytes represent the same value (or they are both invalid).
     /// This is the "guaranteed" padding. There may be more bytes that are padding for some
     /// but not all variants of this type; those are not included.
     /// (E.g. `Option<i8>` has no guaranteed padding so the empty range set is returned, but its `None` value still has padding).
-    pub fn padding_ranges<C>(&self, cx: &C) -> Vec<Range<Size>>
+    pub fn variant_independent_padding_ranges<C>(&self, cx: &C) -> Vec<Range<Size>>
     where
         Ty: TyAbiInterface<'a, C> + Copy,
     {
@@ -316,6 +337,45 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
         uninit_ranges
     }
 
+    /// The ranges of bytes that are ignored by the representation relation of this variant.
+    ///
+    /// The result does not include variant-independent padding.
+    pub fn variant_dependent_padding_ranges<C>(
+        &self,
+        cx: &C,
+        variant_index: VariantIdx,
+    ) -> Vec<Range<Size>>
+    where
+        Ty: TyAbiInterface<'a, C> + Copy,
+    {
+        let Variants::Multiple { .. } = self.variants else {
+            return Vec::new();
+        };
+
+        // Bytes that are data in some variant.
+        let mut any = RangeSet::new();
+        self.add_data_ranges(cx, Size::ZERO, &mut any);
+
+        // Bytes that are data in this variant.
+        let mut this = RangeSet::new();
+
+        // The variants do not contain e.g. the discriminant or coroutine upvars.
+        let FieldsShape::Arbitrary { offsets, in_memory_order: _ } = &self.fields else {
+            unreachable!("a multi-variant layout should have `Arbitrary` fields")
+        };
+
+        // So add them explicitly.
+        for (field, &offset) in offsets.iter_enumerated() {
+            let field = self.field(cx, field.as_usize());
+            field.add_data_ranges(cx, offset, &mut this);
+        }
+
+        self.for_variant(cx, variant_index).add_data_ranges(cx, Size::ZERO, &mut this);
+
+        // Padding specific to this variant: data in some variant, but not in this one.
+        any.difference(&this).0.iter().map(|&(offset, size)| offset..offset + size).collect()
+    }
+
     /// Extend `out` with all ranges of bytes that *may* carry relevant data for values of this type.
     /// For enums and unions there are offsets that are initialized for some
     /// variants but not for others; those offset *will* get added to `out`.
@@ -327,39 +387,42 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
             return;
         }
 
-        match &self.variants {
-            Variants::Empty => { /* done */ }
-            Variants::Single { index: _ } => match &self.fields {
-                FieldsShape::Primitive => {
-                    out.add_range(base_offset, self.size);
-                }
-                &FieldsShape::Union(field_count) => {
-                    for field in 0..field_count.get() {
-                        let field = self.field(cx, field);
-                        field.add_data_ranges(cx, base_offset, out);
-                    }
+        // Visit the fields of this value. For enum values the fields include the discriminant.
+        match &self.fields {
+            FieldsShape::Primitive => {
+                out.add_range(base_offset, self.size);
+            }
+            &FieldsShape::Union(field_count) => {
+                for field in 0..field_count.get() {
+                    let field = self.field(cx, field);
+                    field.add_data_ranges(cx, base_offset, out);
                 }
-                &FieldsShape::Array { stride, count } => {
-                    let elem = self.field(cx, 0);
-
-                    // For scalars we know there is no padding between the elements,
-                    // so the entire array is a single big data range.
-                    if elem.backend_repr.is_scalar() {
-                        out.add_range(base_offset, elem.size * count);
-                    } else {
-                        // FIXME: this is really inefficient for large arrays.
-                        for idx in 0..count {
-                            elem.add_data_ranges(cx, base_offset + idx * stride, out);
-                        }
+            }
+            &FieldsShape::Array { stride, count } => {
+                let elem = self.field(cx, 0);
+
+                // For scalars we know there is no padding between the elements,
+                // so the entire array is a single big data range.
+                if elem.backend_repr.is_scalar() {
+                    out.add_range(base_offset, elem.size * count);
+                } else {
+                    // FIXME: this is really inefficient for large arrays.
+                    for idx in 0..count {
+                        elem.add_data_ranges(cx, base_offset + idx * stride, out);
                     }
                 }
-                FieldsShape::Arbitrary { offsets, in_memory_order: _ } => {
-                    for (field, &offset) in offsets.iter_enumerated() {
-                        let field = self.field(cx, field.as_usize());
-                        field.add_data_ranges(cx, base_offset + offset, out);
-                    }
+            }
+            FieldsShape::Arbitrary { offsets, in_memory_order: _ } => {
+                for (field, &offset) in offsets.iter_enumerated() {
+                    let field = self.field(cx, field.as_usize());
+                    field.add_data_ranges(cx, base_offset + offset, out);
                 }
-            },
+            }
+        }
+
+        // Visit the fields of each variant.
+        match &self.variants {
+            Variants::Empty | Variants::Single { index: _ } => { /* done */ }
             Variants::Multiple { variants, .. } => {
                 for variant in variants.indices() {
                     let variant = self.for_variant(cx, variant);

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -2,7 +2,8 @@ use std::cmp;
 use std::ops::Range;
 
 use rustc_abi::{
-    Align, ArmCall, BackendRepr, CanonAbi, ExternAbi, HasDataLayout, Reg, Size, WrappingRange,
+    Align, ArmCall, BackendRepr, CanonAbi, ExternAbi, FieldsShape, HasDataLayout, Reg, Size,
+    VariantIdx, Variants, WrappingRange,
 };
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
@@ -11,7 +12,7 @@ use rustc_hir::attrs::AttributeKind;
 use rustc_hir::lang_items::LangItem;
 use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER;
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
-use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
+use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout, ValidityRequirement};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
 use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt};
 use rustc_middle::{bug, span_bug};
@@ -603,15 +604,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                 if self.fn_abi.conv == CanonAbi::Arm(ArmCall::CCmseNonSecureEntry) {
                     // The return value of an `extern "cmse-nonsecure-entry"` function crosses the
-                    // secure boundary. Zero padding bytes so information does not leak.
-                    //
-                    // This only zeroes "guaranteed" padding. There may be more bytes that are
-                    // padding for some but not all variants of this type; those are not zeroed.
-                    //
-                    // Returning a value with value-dependent padding will instead trigger a lint.
+                    // secure boundary. Clear any padding bytes so information does not leak.
                     let ret_layout = self.fn_abi.ret.layout;
-                    let uninit_ranges = ret_layout.padding_ranges(bx.cx());
-                    self.zero_byte_ranges(bx, llslot, ret_layout.size, &uninit_ranges);
+                    self.clear_padding_cmse(bx, llslot, ret_layout.size, ret_layout);
                 }
 
                 load_cast(bx, cast_ty, llslot, self.fn_abi.ret.layout.align.abi)
@@ -1730,22 +1725,166 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
     }
 
+    /// When using CMSE, values that cross the secure boundary from secure to non-secure mode can
+    /// contain stale secure data in their padding bytes. This function clears that data. This is
+    /// required when a value is:
+    ///
+    /// - passed to an `extern "cmse-nonsecure-call"` function
+    /// - returned from an `extern "cmse-nonsecure-entry"` function
+    ///
+    /// This function clears both:
+    ///
+    /// - variant-independent padding, bytes that are padding for all valid values of the type
+    /// - variant-dependent padding, bytes that are padding for some but not all values of the type
+    ///
+    /// Clearing variant-dependent padding requires looking at the data at runtime to determine what
+    /// bytes to clear.
+    fn clear_padding_cmse(
+        &mut self,
+        bx: &mut Bx,
+        base_ptr: Bx::Value,
+        limit: Size,
+        layout: TyAndLayout<'tcx>,
+    ) {
+        // First clear variant-independent padding, a series of memsets.
+        let variant_independent = layout.variant_independent_padding_ranges(self.cx);
+        self.zero_byte_ranges(bx, base_ptr, Size::ZERO, limit, &variant_independent);
+
+        // Then clear the extra padding of the active variant of any (nested) enum.
+        self.clear_variant_dependent_padding(bx, base_ptr, Size::ZERO, limit, layout);
+    }
+
+    fn clear_variant_dependent_padding(
+        &mut self,
+        bx: &mut Bx,
+        base_ptr: Bx::Value,
+        base_offset: Size,
+        limit: Size,
+        layout: TyAndLayout<'tcx>,
+    ) {
+        let cx = self.cx;
+
+        if !layout.has_variant_dependent_padding(cx) {
+            return;
+        }
+
+        // Recurse into aggregate fields/elements to reach any nested enums.
+        match layout.fields {
+            FieldsShape::Array { stride, count } => {
+                let elem = layout.field(cx, 0);
+                if elem.has_variant_dependent_padding(cx) {
+                    for idx in 0..count {
+                        let off = base_offset + idx * stride;
+                        self.clear_variant_dependent_padding(bx, base_ptr, off, limit, elem);
+                    }
+                }
+            }
+            FieldsShape::Arbitrary { .. } => {
+                for i in 0..layout.fields.count() {
+                    let field = layout.field(cx, i);
+                    if field.has_variant_dependent_padding(cx) {
+                        let off = base_offset + layout.fields.offset(i);
+                        self.clear_variant_dependent_padding(bx, base_ptr, off, limit, field);
+                    }
+                }
+            }
+            FieldsShape::Primitive | FieldsShape::Union(_) => { /* nothing to visit */ }
+        }
+
+        // If this is not a multi-variant enum, we're done.
+        let Variants::Multiple { ref variants, .. } = layout.variants else {
+            return;
+        };
+
+        // Collect variants that will need padding cleared.
+        let mut work = Vec::with_capacity(variants.len());
+        for i in 0..variants.len() {
+            let idx = VariantIdx::from_usize(i);
+            let variant = layout.for_variant(cx, idx);
+
+            // Don't consider uninhabited variants.
+            if variant.is_uninhabited() {
+                continue;
+            }
+
+            let variant_dependent = layout.variant_dependent_padding_ranges(cx, idx);
+            let has_nested_variant_dependent = (0..variant.fields.count())
+                .any(|i| variant.field(cx, i).has_variant_dependent_padding(cx));
+
+            if !variant_dependent.is_empty() || has_nested_variant_dependent {
+                work.push((idx, variant, variant_dependent));
+            }
+        }
+
+        if work.is_empty() {
+            return;
+        }
+
+        // Build the switch and clear the appropriate padding for each variant.
+        let root_block = bx.llbb();
+        let join_block = bx.append_sibling_block("cmse_pad_join");
+        let mut cases = Vec::with_capacity(work.len());
+
+        for (idx, variant, variant_dependent) in work.into_iter() {
+            let Some(discr) = layout.ty.discriminant_for_variant(bx.tcx(), idx) else {
+                bug!("multi-variant layout on a type without discriminants");
+            };
+
+            let variant_block = bx.append_sibling_block("cmse_pad_variant");
+            bx.switch_to_block(variant_block);
+
+            // Clear the padding of this variant.
+            self.zero_byte_ranges(bx, base_ptr, base_offset, limit, &variant_dependent);
+
+            // Recurse into the fields.
+            for i in 0..variant.fields.count() {
+                let field = variant.field(cx, i);
+                let off = base_offset + variant.fields.offset(i);
+                self.clear_variant_dependent_padding(bx, base_ptr, off, limit, field);
+            }
+
+            bx.br(join_block);
+            cases.push((discr.val, variant_block));
+        }
+
+        // Construct the dispatch.
+        bx.switch_to_block(root_block);
+
+        let discr_ty = layout.ty.discriminant_ty(bx.tcx());
+        let enum_ptr = bx.inbounds_ptradd(base_ptr, bx.const_usize(base_offset.bytes()));
+        let operand = OperandRef {
+            val: OperandValue::Ref(PlaceValue::new_sized(enum_ptr, layout.align.abi)),
+            layout,
+            move_annotation: None,
+        };
+        let discr = operand.codegen_get_discr(self, bx, discr_ty);
+
+        // Default to the join block (for variants without variant-dependent padding).
+        bx.switch(discr, join_block, cases.into_iter());
+
+        bx.switch_to_block(join_block);
+    }
+
     fn zero_byte_ranges(
         &mut self,
         bx: &mut Bx,
         ptr: Bx::Value,
+        offset: Size,
         limit: Size,
         ranges: &[Range<Size>],
     ) {
         let zero = bx.const_u8(0);
 
         for range in ranges {
-            let end = cmp::min(range.end, limit);
+            let start = range.start + offset;
+            let end = range.end + offset;
+
+            let end = cmp::min(end, limit);
             if range.start >= end {
                 continue;
             }
-            let offset = bx.const_usize(range.start.bytes());
-            let len = bx.const_usize((end - range.start).bytes());
+            let offset = bx.const_usize(start.bytes());
+            let len = bx.const_usize((end - start).bytes());
             let ptr = bx.inbounds_ptradd(ptr, offset);
             bx.memset(ptr, zero, len, Align::ONE, MemFlags::empty());
         }
@@ -1880,18 +2019,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 );
 
                 // The arguments of an `extern "cmse-nonsecure-call"` function cross the secure
-                // boundary. Zero padding bytes so information does not leak.
-                //
-                // This only zeroes "guaranteed" padding. There may be more bytes that are
-                // padding for some but not all variants of this type; those are not zeroed.
-                //
-                // Passing an argument with value-dependent padding will instead trigger a lint.
+                // boundary. Clear any padding bytes so information does not leak.
                 if conv == CanonAbi::Arm(ArmCall::CCmseNonSecureCall) {
-                    self.zero_byte_ranges(
+                    self.clear_padding_cmse(
                         bx,
                         llscratch,
                         Size::from_bytes(copy_bytes),
-                        &arg.layout.padding_ranges(bx.cx()),
+                        arg.layout,
                     );
                 }