From 57078384809fffafac4e90e18cc37a91a1dd5200 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Sun, 31 May 2020 09:52:51 +0200
Subject: [PATCH] validate basic sanity for TerminatorKind

---
 src/librustc_mir/interpret/terminator.rs |   8 +-
 src/librustc_mir/transform/validate.rs   | 103 ++++++++++++++++++++++-
 2 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/src/librustc_mir/interpret/terminator.rs b/src/librustc_mir/interpret/terminator.rs
index b048240ca8dc1..3db16a71bab15 100644
--- a/src/librustc_mir/interpret/terminator.rs
+++ b/src/librustc_mir/interpret/terminator.rs
@@ -50,7 +50,13 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.go_to_block(target_block);
             }
 
-            Call { ref func, ref args, destination, ref cleanup, .. } => {
+            Call {
+                ref func,
+                ref args,
+                destination,
+                ref cleanup,
+                from_hir_call: _from_hir_call,
+            } => {
                 let old_stack = self.frame_idx();
                 let old_loc = self.frame().loc;
                 let func = self.eval_operand(func, None)?;
diff --git a/src/librustc_mir/transform/validate.rs b/src/librustc_mir/transform/validate.rs
index a25edd131baa1..046889193dac3 100644
--- a/src/librustc_mir/transform/validate.rs
+++ b/src/librustc_mir/transform/validate.rs
@@ -3,8 +3,11 @@
 use super::{MirPass, MirSource};
 use rustc_middle::mir::visit::Visitor;
 use rustc_middle::{
-    mir::{Body, Location, Operand, Rvalue, Statement, StatementKind},
-    ty::{ParamEnv, TyCtxt},
+    mir::{
+        BasicBlock, Body, Location, Operand, Rvalue, Statement, StatementKind, Terminator,
+        TerminatorKind,
+    },
+    ty::{self, ParamEnv, TyCtxt},
 };
 use rustc_span::{def_id::DefId, Span, DUMMY_SP};
 
@@ -38,6 +41,12 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             &format!("broken MIR in {:?} ({}): {}", self.def_id, self.when, msg.as_ref()),
         );
     }
+
+    fn check_bb(&self, span: Span, bb: BasicBlock) {
+        if self.body.basic_blocks().get(bb).is_none() {
+            self.fail(span, format!("encountered jump to invalid basic block {:?}", bb))
+        }
+    }
 }
 
 impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
@@ -77,4 +86,94 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
             }
         }
     }
+
+    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, _location: Location) {
+        match &terminator.kind {
+            TerminatorKind::Goto { target } => {
+                self.check_bb(terminator.source_info.span, *target);
+            }
+            TerminatorKind::SwitchInt { targets, .. } => {
+                if targets.is_empty() {
+                    self.fail(
+                        terminator.source_info.span,
+                        "encountered `SwitchInt` terminator with no target to jump to",
+                    );
+                }
+                for target in targets {
+                    self.check_bb(terminator.source_info.span, *target);
+                }
+            }
+            TerminatorKind::Drop { target, unwind, .. } => {
+                self.check_bb(terminator.source_info.span, *target);
+                if let Some(unwind) = unwind {
+                    self.check_bb(terminator.source_info.span, *unwind);
+                }
+            }
+            TerminatorKind::DropAndReplace { target, unwind, .. } => {
+                self.check_bb(terminator.source_info.span, *target);
+                if let Some(unwind) = unwind {
+                    self.check_bb(terminator.source_info.span, *unwind);
+                }
+            }
+            TerminatorKind::Call { func, destination, cleanup, .. } => {
+                let func_ty = func.ty(&self.body.local_decls, self.tcx);
+                match func_ty.kind {
+                    ty::FnPtr(..) | ty::FnDef(..) => {}
+                    _ => self.fail(
+                        terminator.source_info.span,
+                        format!("encountered non-callable type {} in `Call` terminator", func_ty),
+                    ),
+                }
+                if let Some((_, target)) = destination {
+                    self.check_bb(terminator.source_info.span, *target);
+                }
+                if let Some(cleanup) = cleanup {
+                    self.check_bb(terminator.source_info.span, *cleanup);
+                }
+            }
+            TerminatorKind::Assert { cond, target, cleanup, .. } => {
+                let cond_ty = cond.ty(&self.body.local_decls, self.tcx);
+                if cond_ty != self.tcx.types.bool {
+                    self.fail(
+                        terminator.source_info.span,
+                        format!(
+                            "encountered non-boolean condition of type {} in `Assert` terminator",
+                            cond_ty
+                        ),
+                    );
+                }
+                self.check_bb(terminator.source_info.span, *target);
+                if let Some(cleanup) = cleanup {
+                    self.check_bb(terminator.source_info.span, *cleanup);
+                }
+            }
+            TerminatorKind::Yield { resume, drop, .. } => {
+                self.check_bb(terminator.source_info.span, *resume);
+                if let Some(drop) = drop {
+                    self.check_bb(terminator.source_info.span, *drop);
+                }
+            }
+            TerminatorKind::FalseEdges { real_target, imaginary_target } => {
+                self.check_bb(terminator.source_info.span, *real_target);
+                self.check_bb(terminator.source_info.span, *imaginary_target);
+            }
+            TerminatorKind::FalseUnwind { real_target, unwind } => {
+                self.check_bb(terminator.source_info.span, *real_target);
+                if let Some(unwind) = unwind {
+                    self.check_bb(terminator.source_info.span, *unwind);
+                }
+            }
+            TerminatorKind::InlineAsm { destination, .. } => {
+                if let Some(destination) = destination {
+                    self.check_bb(terminator.source_info.span, *destination);
+                }
+            }
+            // Nothing to validate for these.
+            TerminatorKind::Resume
+            | TerminatorKind::Abort
+            | TerminatorKind::Return
+            | TerminatorKind::Unreachable
+            | TerminatorKind::GeneratorDrop => {}
+        }
+    }
 }