validate enum discriminant whenever it is read

8 files changed, 64 insertions, 65 deletions

diff --git a/src/librustc/mir/interpret/error.rs b/src/librustc/mir/interpret/error.rs
index dc6d17d3453..c62ed841866 100644
--- a/src/librustc/mir/interpret/error.rs
+++ b/src/librustc/mir/interpret/error.rs
@@ -303,7 +303,7 @@ impl<'tcx, O> EvalErrorKind<'tcx, O> {
             InvalidBool =>
                 "invalid boolean value read",
             InvalidDiscriminant =>
-                "invalid enum discriminant value read",
+                "invalid enum discriminant value read or written",
             PointerOutOfBounds { .. } =>
                 "pointer offset outside bounds of allocation",
             InvalidNullPointerUsage =>
diff --git a/src/librustc_mir/const_eval.rs b/src/librustc_mir/const_eval.rs
index 3f0ab3b8f7c..3017197477c 100644
--- a/src/librustc_mir/const_eval.rs
+++ b/src/librustc_mir/const_eval.rs
@@ -382,7 +382,7 @@ pub fn const_variant_index<'a, 'tcx>(
     trace!("const_variant_index: {:?}, {:?}", instance, val);
     let ecx = mk_eval_cx(tcx, instance, param_env).unwrap();
     let op = ecx.const_to_op(val)?;
-    ecx.read_discriminant_as_variant_index(op)
+    Ok(ecx.read_discriminant(op)?.1)
 }
 
 pub fn const_to_allocation_provider<'a, 'tcx>(
diff --git a/src/librustc_mir/interpret/operand.rs b/src/librustc_mir/interpret/operand.rs
index 6f3fc106d26..32ea6351406 100644
--- a/src/librustc_mir/interpret/operand.rs
+++ b/src/librustc_mir/interpret/operand.rs
@@ -544,34 +544,11 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
         self.const_value_to_op(cv.val)
     }
 
-    /// reads a tag and produces the corresponding variant index
-    pub fn read_discriminant_as_variant_index(
+    /// Read discriminant, return the runtime value as well as the variant index.
+    pub fn read_discriminant(
         &self,
         rval: OpTy<'tcx>,
-    ) -> EvalResult<'tcx, usize> {
-        match rval.layout.variants {
-            layout::Variants::Single { index } => Ok(index),
-            layout::Variants::Tagged { .. } => {
-                let discr_val = self.read_discriminant_value(rval)?;
-                rval.layout.ty
-                    .ty_adt_def()
-                    .expect("tagged layout for non adt")
-                    .discriminants(self.tcx.tcx)
-                    .position(|var| var.val == discr_val)
-                    .ok_or_else(|| EvalErrorKind::InvalidDiscriminant.into())
-            }
-            layout::Variants::NicheFilling { .. } => {
-                let discr_val = self.read_discriminant_value(rval)?;
-                assert_eq!(discr_val as usize as u128, discr_val);
-                Ok(discr_val as usize)
-            },
-        }
-    }
-
-    pub fn read_discriminant_value(
-        &self,
-        rval: OpTy<'tcx>,
-    ) -> EvalResult<'tcx, u128> {
+    ) -> EvalResult<'tcx, (u128, usize)> {
         trace!("read_discriminant_value {:#?}", rval.layout);
         if rval.layout.abi == layout::Abi::Uninhabited {
             return err!(Unreachable);
@@ -582,20 +559,21 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                 let discr_val = rval.layout.ty.ty_adt_def().map_or(
                     index as u128,
                     |def| def.discriminant_for_variant(*self.tcx, index).val);
-                return Ok(discr_val);
+                return Ok((discr_val, index));
             }
             layout::Variants::Tagged { .. } |
             layout::Variants::NicheFilling { .. } => {},
         }
+        // read raw discriminant value
         let discr_op = self.operand_field(rval, 0)?;
         let discr_val = self.read_value(discr_op)?;
-        trace!("discr value: {:?}", discr_val);
         let raw_discr = discr_val.to_scalar()?;
+        trace!("discr value: {:?}", raw_discr);
+        // post-process
         Ok(match rval.layout.variants {
             layout::Variants::Single { .. } => bug!(),
-            // FIXME: We should catch invalid discriminants here!
             layout::Variants::Tagged { .. } => {
-                if discr_val.layout.ty.is_signed() {
+                let real_discr = if discr_val.layout.ty.is_signed() {
                     let i = raw_discr.to_bits(discr_val.layout.size)? as i128;
                     // going from layout tag type to typeck discriminant type
                     // requires first sign extending with the layout discriminant
@@ -612,7 +590,15 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                     (truncatee << shift) >> shift
                 } else {
                     raw_discr.to_bits(discr_val.layout.size)?
-                }
+                };
+                // Make sure we catch invalid discriminants
+                let index = rval.layout.ty
+                    .ty_adt_def()
+                    .expect("tagged layout for non adt")
+                    .discriminants(self.tcx.tcx)
+                    .position(|var| var.val == real_discr)
+                    .ok_or_else(|| EvalErrorKind::InvalidDiscriminant)?;
+                (real_discr, index)
             },
             layout::Variants::NicheFilling {
                 dataful_variant,
@@ -622,8 +608,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
             } => {
                 let variants_start = *niche_variants.start() as u128;
                 let variants_end = *niche_variants.end() as u128;
-                match raw_discr {
+                let real_discr = match raw_discr {
                     Scalar::Ptr(_) => {
+                        // The niche must be just 0 (which a pointer value never is)
                         assert!(niche_start == 0);
                         assert!(variants_start == variants_end);
                         dataful_variant as u128
@@ -638,7 +625,14 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                             dataful_variant as u128
                         }
                     },
-                }
+                };
+                let index = real_discr as usize;
+                assert_eq!(index as u128, real_discr);
+                assert!(index < rval.layout.ty
+                    .ty_adt_def()
+                    .expect("tagged layout for non adt")
+                    .variants.len());
+                (real_discr, index)
             }
         })
     }
diff --git a/src/librustc_mir/interpret/place.rs b/src/librustc_mir/interpret/place.rs
index 6d8d1a1e8ff..b1eb3749d5d 100644
--- a/src/librustc_mir/interpret/place.rs
+++ b/src/librustc_mir/interpret/place.rs
@@ -702,7 +702,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
         Ok(MPlaceTy::from_aligned_ptr(ptr, layout))
     }
 
-    pub fn write_discriminant_value(
+    pub fn write_discriminant_index(
         &mut self,
         variant_index: usize,
         dest: PlaceTy<'tcx>,
@@ -710,14 +710,15 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
         match dest.layout.variants {
             layout::Variants::Single { index } => {
                 if index != variant_index {
-                    // If the layout of an enum is `Single`, all
-                    // other variants are necessarily uninhabited.
-                    assert_eq!(dest.layout.for_variant(&self, variant_index).abi,
-                               layout::Abi::Uninhabited);
+                    return err!(InvalidDiscriminant);
                 }
             }
             layout::Variants::Tagged { ref tag, .. } => {
-                let discr_val = dest.layout.ty.ty_adt_def().unwrap()
+                let adt_def = dest.layout.ty.ty_adt_def().unwrap();
+                if variant_index >= adt_def.variants.len() {
+                    return err!(InvalidDiscriminant);
+                }
+                let discr_val = adt_def
                     .discriminant_for_variant(*self.tcx, variant_index)
                     .val;
 
@@ -740,6 +741,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                 niche_start,
                 ..
             } => {
+                if variant_index >= dest.layout.ty.ty_adt_def().unwrap().variants.len() {
+                    return err!(InvalidDiscriminant);
+                }
                 if variant_index != dataful_variant {
                     let niche_dest =
                         self.place_field(dest, 0)?;
diff --git a/src/librustc_mir/interpret/step.rs b/src/librustc_mir/interpret/step.rs
index 54fdf8e0d4b..b2faf59af3e 100644
--- a/src/librustc_mir/interpret/step.rs
+++ b/src/librustc_mir/interpret/step.rs
@@ -127,7 +127,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                 variant_index,
             } => {
                 let dest = self.eval_place(place)?;
-                self.write_discriminant_value(variant_index, dest)?;
+                self.write_discriminant_index(variant_index, dest)?;
             }
 
             // Mark locals as alive
@@ -222,7 +222,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
             Aggregate(ref kind, ref operands) => {
                 let (dest, active_field_index) = match **kind {
                     mir::AggregateKind::Adt(adt_def, variant_index, _, _, active_field_index) => {
-                        self.write_discriminant_value(variant_index, dest)?;
+                        self.write_discriminant_index(variant_index, dest)?;
                         if adt_def.is_enum() {
                             (self.place_downcast(dest, variant_index)?, active_field_index)
                         } else {
@@ -312,7 +312,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
 
             Discriminant(ref place) => {
                 let place = self.eval_place(place)?;
-                let discr_val = self.read_discriminant_value(self.place_to_op(place)?)?;
+                let discr_val = self.read_discriminant(self.place_to_op(place)?)?.0;
                 let size = dest.layout.size.bytes() as u8;
                 self.write_scalar(Scalar::Bits {
                     bits: discr_val,
diff --git a/src/librustc_mir/interpret/validity.rs b/src/librustc_mir/interpret/validity.rs
index af191ce3adc..a3752a05fc8 100644
--- a/src/librustc_mir/interpret/validity.rs
+++ b/src/librustc_mir/interpret/validity.rs
@@ -198,25 +198,25 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
         seen: &mut FxHashSet<(OpTy<'tcx>)>,
         todo: &mut Vec<(OpTy<'tcx>, Vec<PathElem>)>,
     ) -> EvalResult<'tcx> {
-        trace!("validate_mplace: {:?}, {:#?}", *dest, dest.layout);
+        trace!("validate_operand: {:?}, {:#?}", *dest, dest.layout);
 
         // Find the right variant.  We have to handle this as a prelude, not via
         // proper recursion with the new inner layout, to be able to later nicely
         // print the field names of the enum field that is being accessed.
         let (variant, dest) = match dest.layout.variants {
-            layout::Variants::NicheFilling { niche: ref tag, .. } |
-            layout::Variants::Tagged { ref tag, .. } => {
-                let size = tag.value.size(self);
-                // we first read the tag value as scalar, to be able to validate it
-                let tag_mplace = self.operand_field(dest, 0)?;
-                let tag_value = self.read_scalar(tag_mplace)?;
-                path.push(PathElem::Tag);
-                self.validate_scalar(
-                    tag_value, size, tag, &path, tag_mplace.layout.ty
-                )?;
-                path.pop(); // remove the element again
-                // then we read it again to get the index, to continue
-                let variant = self.read_discriminant_as_variant_index(dest.into())?;
+            layout::Variants::NicheFilling { .. } |
+            layout::Variants::Tagged { .. } => {
+                let variant = match self.read_discriminant(dest) {
+                    Ok(res) => res.1,
+                    Err(err) => match err.kind {
+                        EvalErrorKind::InvalidDiscriminant |
+                        EvalErrorKind::ReadPointerAsBytes =>
+                            return validation_failure!(
+                                "invalid enum discriminant", path
+                            ),
+                        _ => return Err(err),
+                    }
+                };
                 let inner_dest = self.operand_downcast(dest, variant)?;
                 // Put the variant projection onto the path, as a field
                 path.push(PathElem::Field(dest.layout.ty
@@ -258,17 +258,17 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                                 let alloc_kind = self.tcx.alloc_map.lock().get(ptr.alloc_id);
                                 if let Some(AllocType::Static(did)) = alloc_kind {
                                     // statics from other crates are already checked.
-                                    // extern statics should not be validated as they have no body.
+                                    // extern statics cannot be validated as they have no body.
                                     if !did.is_local() || self.tcx.is_foreign_item(did) {
                                         return Ok(());
                                     }
                                 }
                                 if value.layout.ty.builtin_deref(false).is_some() {
-                                    trace!("Recursing below ptr {:#?}", value);
                                     let ptr_op = self.ref_to_mplace(value)?.into();
                                     // we have not encountered this pointer+layout combination
                                     // before.
                                     if seen.insert(ptr_op) {
+                                        trace!("Recursing below ptr {:#?}", *value);
                                         todo.push((ptr_op, path_clone_and_deref(path)));
                                     }
                                 }
@@ -284,6 +284,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                 // See https://github.com/rust-lang/rust/issues/32836#issuecomment-406875389
             },
             layout::FieldPlacement::Array { .. } => {
+                // FIXME: For a TyStr, check that this is valid UTF-8.
                 // Skips for ZSTs; we could have an empty array as an immediate
                 if !dest.layout.is_zst() {
                     let dest = dest.to_mem_place(); // arrays cannot be immediate
@@ -316,8 +317,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                     let unpacked_ptr = self.unpack_unsized_mplace(ptr)?.into();
                     // for safe ptrs, recursively check it
                     if !dest.layout.ty.is_unsafe_ptr() {
-                        trace!("Recursing below fat ptr {:?} (unpacked: {:?})", ptr, unpacked_ptr);
                         if seen.insert(unpacked_ptr) {
+                            trace!("Recursing below fat ptr {:?} (unpacked: {:?})",
+                                ptr, unpacked_ptr);
                             todo.push((unpacked_ptr, path_clone_and_deref(path)));
                         }
                     }
@@ -329,7 +331,6 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
                         self.validate_operand(field, path, seen, todo)?;
                         path.truncate(path_len);
                     }
-                    // FIXME: For a TyStr, check that this is valid UTF-8.
                 }
             }
         }
diff --git a/src/test/ui/consts/const-eval/double_check2.stderr b/src/test/ui/consts/const-eval/double_check2.stderr
index 739af12d09c..6b33007ec09 100644
--- a/src/test/ui/consts/const-eval/double_check2.stderr
+++ b/src/test/ui/consts/const-eval/double_check2.stderr
@@ -5,7 +5,7 @@ LL | / static FOO: (&Foo, &Bar) = unsafe {( //~ undefined behavior
 LL | |     Union { usize: &BAR }.foo,
 LL | |     Union { usize: &BAR }.bar,
 LL | | )};
-   | |___^ type validation failed: encountered 5 at .1.<deref>.<enum-tag>, but expected something in the range 42..=99
+   | |___^ type validation failed: encountered invalid enum discriminant at .1.<deref>
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rust compiler repository if you believe it should not be considered undefined behavior
 
diff --git a/src/test/ui/consts/const-eval/ub-enum.stderr b/src/test/ui/consts/const-eval/ub-enum.stderr
index 98e9b598b54..bb015db3e54 100644
--- a/src/test/ui/consts/const-eval/ub-enum.stderr
+++ b/src/test/ui/consts/const-eval/ub-enum.stderr
@@ -2,7 +2,7 @@ error[E0080]: this constant likely exhibits undefined behavior
   --> $DIR/ub-enum.rs:22:1
    |
 LL | const BAD_ENUM: Enum = unsafe { TransmuteEnum { a: &1 }.b };
-   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed: encountered pointer at .<enum-tag>, but expected something in the range 0..=0
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed: encountered invalid enum discriminant
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rust compiler repository if you believe it should not be considered undefined behavior
 
@@ -10,7 +10,7 @@ error[E0080]: this constant likely exhibits undefined behavior
   --> $DIR/ub-enum.rs:35:1
    |
 LL | const BAD_ENUM2 : Enum2 = unsafe { TransmuteEnum2 { a: 0 }.b };
-   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed: encountered 0 at .<enum-tag>, but expected something in the range 2..=2
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed: encountered invalid enum discriminant
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rust compiler repository if you believe it should not be considered undefined behavior