| use rustc_abi::{Scalar, Size, TagEncoding, Variants, WrappingRange}; | |
| use rustc_data_structures::thin_vec::ThinVec; | |
| use rustc_hir::LangItem; | |
| use rustc_index::IndexVec; | |
| use rustc_middle::bug; | |
| use rustc_middle::mir::visit::Visitor; | |
| use rustc_middle::mir::*; | |
| use rustc_middle::ty::layout::PrimitiveExt; | |
| use rustc_middle::ty::{self, Ty, TyCtxt, TypingEnv}; | |
| use rustc_session::Session; | |
| use tracing::debug; | |
| use crate::PassPolicy; | |
| /// This pass inserts checks for a valid enum discriminant where they are most | |
| /// likely to find UB, because checking everywhere like Miri would generate too | |
| /// much MIR. | |
| pub(super) struct CheckEnums; | |
| impl<'tcx> crate::MirPass<'tcx> for CheckEnums { | |
| fn policy(&self, sess: &Session) -> PassPolicy { | |
| // When UB checks are enabled this is part of their semantics, not an optimization. | |
| PassPolicy::optional_non_optimization(sess.ub_checks()) | |
| } | |
| fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { | |
| // This pass emits new panics. If for whatever reason we do not have a panic | |
| // implementation, running this pass may cause otherwise-valid code to not compile. | |
| if tcx.lang_items().get(LangItem::PanicImpl).is_none() { | |
| return; | |
| } | |
| let typing_env = body.typing_env(tcx); | |
| let basic_blocks = body.basic_blocks.as_mut(); | |
| let local_decls = &mut body.local_decls; | |
| // This operation inserts new blocks. Each insertion changes the Location for all | |
| // statements/blocks after. Iterating or visiting the MIR in order would require updating | |
| // our current location after every insertion. By iterating backwards, we dodge this issue: | |
| // The only Locations that an insertion changes have already been handled. | |
| for block in basic_blocks.indices().rev() { | |
| for statement_index in (0..basic_blocks[block].statements.len()).rev() { | |
| let location = Location { block, statement_index }; | |
| let statement = &basic_blocks[block].statements[statement_index]; | |
| let source_info = statement.source_info; | |
| let mut finder = EnumFinder::new(tcx, local_decls, typing_env); | |
| finder.visit_statement(statement, location); | |
| for check in finder.into_found_enums() { | |
| debug!("Inserting enum check"); | |
| let new_block = split_block(basic_blocks, location); | |
| match check { | |
| EnumCheckType::Direct { op_size, .. } | |
| | EnumCheckType::WithNiche { op_size, .. } | |
| if op_size.bytes() == 0 => | |
| { | |
| // It is never valid to use a ZST as a discriminant for an inhabited enum, but that will | |
| // have been caught by the type checker. Do nothing but ensure that a bug has been signaled. | |
| tcx.dcx().span_delayed_bug( | |
| source_info.span, | |
| "cannot build enum discriminant from zero-sized type", | |
| ); | |
| basic_blocks[block].terminator = Some(Terminator { | |
| source_info, | |
| kind: TerminatorKind::Goto { target: new_block }, | |
| attributes: ThinVec::new(), | |
| }); | |
| } | |
| EnumCheckType::Direct { source_op, discr, op_size, valid_discrs } => { | |
| insert_direct_enum_check( | |
| tcx, | |
| local_decls, | |
| basic_blocks, | |
| block, | |
| source_op, | |
| discr, | |
| op_size, | |
| valid_discrs, | |
| source_info, | |
| new_block, | |
| ) | |
| } | |
| EnumCheckType::Uninhabited => insert_uninhabited_enum_check( | |
| tcx, | |
| local_decls, | |
| &mut basic_blocks[block], | |
| source_info, | |
| new_block, | |
| ), | |
| EnumCheckType::WithNiche { | |
| source_op, | |
| discr, | |
| op_size, | |
| offset, | |
| valid_range, | |
| } => insert_niche_check( | |
| tcx, | |
| local_decls, | |
| &mut basic_blocks[block], | |
| source_op, | |
| valid_range, | |
| discr, | |
| op_size, | |
| offset, | |
| source_info, | |
| new_block, | |
| ), | |
| } | |
| } | |
| } | |
| } | |
| } | |
| } | |
| /// Represent the different kind of enum checks we can insert. | |
| enum EnumCheckType<'tcx> { | |
| /// We know we try to create an uninhabited enum from an inhabited variant. | |
| Uninhabited, | |
| /// We know the enum does no niche optimizations and can thus easily compute | |
| /// the valid discriminants. | |
| Direct { | |
| source_op: Operand<'tcx>, | |
| discr: TyAndSize<'tcx>, | |
| op_size: Size, | |
| valid_discrs: Vec<u128>, | |
| }, | |
| /// We try to construct an enum that has a niche. | |
| WithNiche { | |
| source_op: Operand<'tcx>, | |
| discr: TyAndSize<'tcx>, | |
| op_size: Size, | |
| offset: Size, | |
| valid_range: WrappingRange, | |
| }, | |
| } | |
| #[derive(Debug, Copy, Clone)] | |
| struct TyAndSize<'tcx> { | |
| pub ty: Ty<'tcx>, | |
| pub size: Size, | |
| } | |
| /// A [Visitor] that finds the construction of enums and evaluates which checks | |
| /// we should apply. | |
| struct EnumFinder<'a, 'tcx> { | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &'a mut LocalDecls<'tcx>, | |
| typing_env: TypingEnv<'tcx>, | |
| enums: Vec<EnumCheckType<'tcx>>, | |
| } | |
| impl<'a, 'tcx> EnumFinder<'a, 'tcx> { | |
| fn new( | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &'a mut LocalDecls<'tcx>, | |
| typing_env: TypingEnv<'tcx>, | |
| ) -> Self { | |
| EnumFinder { tcx, local_decls, typing_env, enums: Vec::new() } | |
| } | |
| /// Returns the found enum creations and which checks should be inserted. | |
| fn into_found_enums(self) -> Vec<EnumCheckType<'tcx>> { | |
| self.enums | |
| } | |
| } | |
| impl<'a, 'tcx> Visitor<'tcx> for EnumFinder<'a, 'tcx> { | |
| fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) { | |
| if let Rvalue::Cast(CastKind::Transmute, op, ty) = rvalue { | |
| let ty::Adt(adt_def, _) = ty.kind() else { | |
| return; | |
| }; | |
| if !adt_def.is_enum() { | |
| return; | |
| } | |
| let Ok(enum_layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else { | |
| return; | |
| }; | |
| let Ok(op_layout) = self | |
| .tcx | |
| .layout_of(self.typing_env.as_query_input(op.ty(self.local_decls, self.tcx))) | |
| else { | |
| return; | |
| }; | |
| match enum_layout.variants { | |
| Variants::Empty if op_layout.is_uninhabited() => return, | |
| // An empty enum that tries to be constructed from an inhabited value, this | |
| // is never correct. | |
| Variants::Empty => { | |
| // The enum layout is uninhabited but we construct it from sth inhabited. | |
| // This is always UB. | |
| self.enums.push(EnumCheckType::Uninhabited); | |
| } | |
| // Construction of Single value enums is always fine. | |
| Variants::Single { .. } => {} | |
| // Construction of an enum with multiple variants but no niche optimizations. | |
| Variants::Multiple { | |
| tag_encoding: TagEncoding::Direct, | |
| tag: Scalar::Initialized { value, .. }, | |
| .. | |
| } => { | |
| let valid_discrs = | |
| adt_def.discriminants(self.tcx).map(|(_, discr)| discr.val).collect(); | |
| let discr = | |
| TyAndSize { ty: value.to_int_ty(self.tcx), size: value.size(&self.tcx) }; | |
| self.enums.push(EnumCheckType::Direct { | |
| source_op: op.to_copy(), | |
| discr, | |
| op_size: op_layout.size, | |
| valid_discrs, | |
| }); | |
| } | |
| // Construction of an enum with multiple variants and niche optimizations. | |
| Variants::Multiple { | |
| tag_encoding: TagEncoding::Niche { .. }, | |
| tag: Scalar::Initialized { value, valid_range, .. }, | |
| tag_field, | |
| .. | |
| } => { | |
| let discr = | |
| TyAndSize { ty: value.to_int_ty(self.tcx), size: value.size(&self.tcx) }; | |
| self.enums.push(EnumCheckType::WithNiche { | |
| source_op: op.to_copy(), | |
| discr, | |
| op_size: op_layout.size, | |
| offset: enum_layout.fields.offset(tag_field.as_usize()), | |
| valid_range, | |
| }); | |
| } | |
| _ => return, | |
| } | |
| self.super_rvalue(rvalue, location); | |
| } | |
| } | |
| } | |
| fn split_block( | |
| basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'_>>, | |
| location: Location, | |
| ) -> BasicBlock { | |
| let block_data = &mut basic_blocks[location.block]; | |
| // Drain every statement after this one and move the current terminator to a new basic block. | |
| let new_block = BasicBlockData::new_stmts( | |
| block_data.statements.split_off(location.statement_index), | |
| block_data.terminator.take(), | |
| block_data.is_cleanup, | |
| ); | |
| basic_blocks.push(new_block) | |
| } | |
| /// Inserts the cast of an operand (any type) to a u128 value that holds the discriminant value. | |
| fn insert_discr_cast_to_u128<'tcx>( | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, | |
| block_data: &mut BasicBlockData<'tcx>, | |
| source_op: Operand<'tcx>, | |
| discr: TyAndSize<'tcx>, | |
| op_size: Size, | |
| offset: Option<Size>, | |
| source_info: SourceInfo, | |
| ) -> Place<'tcx> { | |
| let get_ty_for_size = |tcx: TyCtxt<'tcx>, size: Size| -> Ty<'tcx> { | |
| match size.bytes() { | |
| 1 => tcx.types.u8, | |
| 2 => tcx.types.u16, | |
| 4 => tcx.types.u32, | |
| 8 => tcx.types.u64, | |
| 16 => tcx.types.u128, | |
| invalid => bug!("Found discriminant with invalid size, has {} bytes", invalid), | |
| } | |
| }; | |
| let (cast_kind, discr_ty_bits) = if discr.size.bytes() < op_size.bytes() { | |
| // The discriminant is less wide than the operand, cast the operand into | |
| // [MaybeUninit; N] and then index into it. | |
| let mu = Ty::new_maybe_uninit(tcx, tcx.types.u8); | |
| let array_len = op_size.bytes(); | |
| let mu_array_ty = Ty::new_array(tcx, mu, array_len); | |
| let mu_array = | |
| local_decls.push(LocalDecl::with_source_info(mu_array_ty, source_info)).into(); | |
| let rvalue = Rvalue::Cast(CastKind::Transmute, source_op, mu_array_ty); | |
| block_data | |
| .statements | |
| .push(Statement::new(source_info, StatementKind::Assign(Box::new((mu_array, rvalue))))); | |
| // Index into the array of MaybeUninit to get something that is actually | |
| // as wide as the discriminant. | |
| let offset = offset.unwrap_or(Size::ZERO); | |
| let smaller_mu_array = mu_array.project_deeper( | |
| &[ProjectionElem::Subslice { | |
| from: offset.bytes(), | |
| to: offset.bytes() + discr.size.bytes(), | |
| from_end: false, | |
| }], | |
| tcx, | |
| ); | |
| (CastKind::Transmute, Operand::Copy(smaller_mu_array)) | |
| } else { | |
| let operand_int_ty = get_ty_for_size(tcx, op_size); | |
| let op_as_int = | |
| local_decls.push(LocalDecl::with_source_info(operand_int_ty, source_info)).into(); | |
| let rvalue = Rvalue::Cast(CastKind::Transmute, source_op, operand_int_ty); | |
| block_data.statements.push(Statement::new( | |
| source_info, | |
| StatementKind::Assign(Box::new((op_as_int, rvalue))), | |
| )); | |
| (CastKind::IntToInt, Operand::Copy(op_as_int)) | |
| }; | |
| // Cast the resulting value to the actual discriminant integer type. | |
| let rvalue = Rvalue::Cast(cast_kind, discr_ty_bits, discr.ty); | |
| let discr_in_discr_ty = | |
| local_decls.push(LocalDecl::with_source_info(discr.ty, source_info)).into(); | |
| block_data.statements.push(Statement::new( | |
| source_info, | |
| StatementKind::Assign(Box::new((discr_in_discr_ty, rvalue))), | |
| )); | |
| // Cast the discriminant to a u128 (base for comparisons of enum discriminants). | |
| let const_u128 = Ty::new_uint(tcx, ty::UintTy::U128); | |
| let rvalue = Rvalue::Cast(CastKind::IntToInt, Operand::Copy(discr_in_discr_ty), const_u128); | |
| let discr = local_decls.push(LocalDecl::with_source_info(const_u128, source_info)).into(); | |
| block_data | |
| .statements | |
| .push(Statement::new(source_info, StatementKind::Assign(Box::new((discr, rvalue))))); | |
| discr | |
| } | |
| fn insert_direct_enum_check<'tcx>( | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, | |
| basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, | |
| current_block: BasicBlock, | |
| source_op: Operand<'tcx>, | |
| discr: TyAndSize<'tcx>, | |
| op_size: Size, | |
| discriminants: Vec<u128>, | |
| source_info: SourceInfo, | |
| new_block: BasicBlock, | |
| ) { | |
| // Insert a new target block that is branched to in case of an invalid discriminant. | |
| let invalid_discr_block_data = BasicBlockData::new(None, false); | |
| let invalid_discr_block = basic_blocks.push(invalid_discr_block_data); | |
| let block_data = &mut basic_blocks[current_block]; | |
| let discr_place = insert_discr_cast_to_u128( | |
| tcx, | |
| local_decls, | |
| block_data, | |
| source_op, | |
| discr, | |
| op_size, | |
| None, | |
| source_info, | |
| ); | |
| // Mask out the bits of the discriminant type. | |
| let mask = discr.size.unsigned_int_max(); | |
| let discr_masked = | |
| local_decls.push(LocalDecl::with_source_info(tcx.types.u128, source_info)).into(); | |
| let rvalue = Rvalue::BinaryOp( | |
| BinOp::BitAnd, | |
| Box::new(( | |
| Operand::Copy(discr_place), | |
| Operand::Constant(Box::new(ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val(ConstValue::from_u128(mask), tcx.types.u128), | |
| })), | |
| )), | |
| ); | |
| block_data | |
| .statements | |
| .push(Statement::new(source_info, StatementKind::Assign(Box::new((discr_masked, rvalue))))); | |
| // Branch based on the discriminant value. | |
| block_data.terminator = Some(Terminator { | |
| source_info, | |
| kind: TerminatorKind::SwitchInt { | |
| discr: Operand::Copy(discr_masked), | |
| targets: SwitchTargets::new( | |
| discriminants | |
| .into_iter() | |
| .map(|discr_val| (discr.size.truncate(discr_val), new_block)), | |
| invalid_discr_block, | |
| ), | |
| }, | |
| attributes: ThinVec::new(), | |
| }); | |
| // Abort in case of an invalid enum discriminant. | |
| basic_blocks[invalid_discr_block].terminator = Some(Terminator { | |
| source_info, | |
| kind: TerminatorKind::Assert { | |
| cond: Operand::Constant(Box::new(ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val(ConstValue::from_bool(false), tcx.types.bool), | |
| })), | |
| expected: true, | |
| target: new_block, | |
| msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr_masked))), | |
| // This calls panic_invalid_enum_construction, which is #[rustc_nounwind]. | |
| // We never want to insert an unwind into unsafe code, because unwinding could | |
| // make a failing UB check turn into much worse UB when we start unwinding. | |
| unwind: UnwindAction::Unreachable, | |
| }, | |
| attributes: ThinVec::new(), | |
| }); | |
| } | |
| fn insert_uninhabited_enum_check<'tcx>( | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, | |
| block_data: &mut BasicBlockData<'tcx>, | |
| source_info: SourceInfo, | |
| new_block: BasicBlock, | |
| ) { | |
| let is_ok: Place<'_> = | |
| local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into(); | |
| block_data.statements.push(Statement::new( | |
| source_info, | |
| StatementKind::Assign(Box::new(( | |
| is_ok, | |
| Rvalue::Use( | |
| Operand::Constant(Box::new(ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val(ConstValue::from_bool(false), tcx.types.bool), | |
| })), | |
| WithRetag::Yes, // it's a bool, retag doesn't matter | |
| ), | |
| ))), | |
| )); | |
| block_data.terminator = Some(Terminator { | |
| source_info, | |
| kind: TerminatorKind::Assert { | |
| cond: Operand::Copy(is_ok), | |
| expected: true, | |
| target: new_block, | |
| msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Constant(Box::new( | |
| ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val(ConstValue::from_u128(0), tcx.types.u128), | |
| }, | |
| )))), | |
| // This calls panic_invalid_enum_construction, which is #[rustc_nounwind]. | |
| // We never want to insert an unwind into unsafe code, because unwinding could | |
| // make a failing UB check turn into much worse UB when we start unwinding. | |
| unwind: UnwindAction::Unreachable, | |
| }, | |
| attributes: ThinVec::new(), | |
| }); | |
| } | |
| fn insert_niche_check<'tcx>( | |
| tcx: TyCtxt<'tcx>, | |
| local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, | |
| block_data: &mut BasicBlockData<'tcx>, | |
| source_op: Operand<'tcx>, | |
| valid_range: WrappingRange, | |
| discr: TyAndSize<'tcx>, | |
| op_size: Size, | |
| offset: Size, | |
| source_info: SourceInfo, | |
| new_block: BasicBlock, | |
| ) { | |
| let discr = insert_discr_cast_to_u128( | |
| tcx, | |
| local_decls, | |
| block_data, | |
| source_op, | |
| discr, | |
| op_size, | |
| Some(offset), | |
| source_info, | |
| ); | |
| // Compare the discriminant against the valid_range. | |
| let start_const = Operand::Constant(Box::new(ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val(ConstValue::from_u128(valid_range.start), tcx.types.u128), | |
| })); | |
| let end_start_diff_const = Operand::Constant(Box::new(ConstOperand { | |
| span: source_info.span, | |
| user_ty: None, | |
| const_: Const::Val( | |
| ConstValue::from_u128(u128::wrapping_sub(valid_range.end, valid_range.start)), | |
| tcx.types.u128, | |
| ), | |
| })); | |
| let discr_diff: Place<'_> = | |
| local_decls.push(LocalDecl::with_source_info(tcx.types.u128, source_info)).into(); | |
| block_data.statements.push(Statement::new( | |
| source_info, | |
| StatementKind::Assign(Box::new(( | |
| discr_diff, | |
| Rvalue::BinaryOp(BinOp::Sub, Box::new((Operand::Copy(discr), start_const))), | |
| ))), | |
| )); | |
| let is_ok: Place<'_> = | |
| local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into(); | |
| block_data.statements.push(Statement::new( | |
| source_info, | |
| StatementKind::Assign(Box::new(( | |
| is_ok, | |
| Rvalue::BinaryOp( | |
| // This is a `WrappingRange`, so make sure to get the wrapping right. | |
| BinOp::Le, | |
| Box::new((Operand::Copy(discr_diff), end_start_diff_const)), | |
| ), | |
| ))), | |
| )); | |
| block_data.terminator = Some(Terminator { | |
| source_info, | |
| kind: TerminatorKind::Assert { | |
| cond: Operand::Copy(is_ok), | |
| expected: true, | |
| target: new_block, | |
| msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr))), | |
| // This calls panic_invalid_enum_construction, which is #[rustc_nounwind]. | |
| // We never want to insert an unwind into unsafe code, because unwinding could | |
| // make a failing UB check turn into much worse UB when we start unwinding. | |
| unwind: UnwindAction::Unreachable, | |
| }, | |
| attributes: ThinVec::new(), | |
| }); | |
| } |