validate enum discriminant whenever it is read

4 files changed, 59 insertions, 60 deletions

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.
                 }
             }
         }