|
| 1 | +use crate::MirPass; |
| 2 | +use rustc_hir::def_id::DefId; |
| 3 | +use rustc_index::vec::IndexVec; |
| 4 | +use rustc_middle::mir::*; |
| 5 | +use rustc_middle::mir::{ |
| 6 | + interpret::{ConstValue, Scalar}, |
| 7 | + visit::{PlaceContext, Visitor}, |
| 8 | +}; |
| 9 | +use rustc_middle::ty::{Ty, TyCtxt, TypeAndMut}; |
| 10 | +use rustc_session::Session; |
| 11 | + |
| 12 | +pub struct CheckAlignment; |
| 13 | + |
| 14 | +impl<'tcx> MirPass<'tcx> for CheckAlignment { |
| 15 | + fn is_enabled(&self, sess: &Session) -> bool { |
| 16 | + sess.opts.debug_assertions |
| 17 | + } |
| 18 | + |
| 19 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 20 | + let basic_blocks = body.basic_blocks.as_mut(); |
| 21 | + let local_decls = &mut body.local_decls; |
| 22 | + |
| 23 | + for block in (0..basic_blocks.len()).rev() { |
| 24 | + let block = block.into(); |
| 25 | + for statement_index in (0..basic_blocks[block].statements.len()).rev() { |
| 26 | + let location = Location { block, statement_index }; |
| 27 | + let statement = &basic_blocks[block].statements[statement_index]; |
| 28 | + let source_info = statement.source_info; |
| 29 | + |
| 30 | + let mut finder = PointerFinder { |
| 31 | + local_decls, |
| 32 | + tcx, |
| 33 | + pointers: Vec::new(), |
| 34 | + def_id: body.source.def_id(), |
| 35 | + }; |
| 36 | + for (pointer, pointee_ty) in finder.find_pointers(statement) { |
| 37 | + debug!("Inserting alignment check for {:?}", pointer.ty(&*local_decls, tcx).ty); |
| 38 | + |
| 39 | + let new_block = split_block(basic_blocks, location); |
| 40 | + insert_alignment_check( |
| 41 | + tcx, |
| 42 | + local_decls, |
| 43 | + &mut basic_blocks[block], |
| 44 | + pointer, |
| 45 | + pointee_ty, |
| 46 | + source_info, |
| 47 | + new_block, |
| 48 | + ); |
| 49 | + } |
| 50 | + } |
| 51 | + } |
| 52 | + } |
| 53 | +} |
| 54 | + |
| 55 | +impl<'tcx, 'a> PointerFinder<'tcx, 'a> { |
| 56 | + fn find_pointers(&mut self, statement: &Statement<'tcx>) -> Vec<(Place<'tcx>, Ty<'tcx>)> { |
| 57 | + self.pointers.clear(); |
| 58 | + self.visit_statement(statement, Location::START); |
| 59 | + core::mem::take(&mut self.pointers) |
| 60 | + } |
| 61 | +} |
| 62 | + |
| 63 | +struct PointerFinder<'tcx, 'a> { |
| 64 | + local_decls: &'a mut LocalDecls<'tcx>, |
| 65 | + tcx: TyCtxt<'tcx>, |
| 66 | + def_id: DefId, |
| 67 | + pointers: Vec<(Place<'tcx>, Ty<'tcx>)>, |
| 68 | +} |
| 69 | + |
| 70 | +impl<'tcx, 'a> Visitor<'tcx> for PointerFinder<'tcx, 'a> { |
| 71 | + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { |
| 72 | + if let PlaceContext::NonUse(_) = context { |
| 73 | + return; |
| 74 | + } |
| 75 | + if !place.is_indirect() { |
| 76 | + return; |
| 77 | + } |
| 78 | + |
| 79 | + let pointer = Place::from(place.local); |
| 80 | + let pointer_ty = pointer.ty(&*self.local_decls, self.tcx).ty; |
| 81 | + |
| 82 | + // We only want to check unsafe pointers |
| 83 | + if !pointer_ty.is_unsafe_ptr() { |
| 84 | + trace!("Indirect, but not an unsafe ptr, not checking {:?}", pointer_ty); |
| 85 | + return; |
| 86 | + } |
| 87 | + |
| 88 | + let Some(pointee) = pointer_ty.builtin_deref(true) else { |
| 89 | + debug!("Indirect but no builtin deref: {:?}", pointer_ty); |
| 90 | + return; |
| 91 | + }; |
| 92 | + let mut pointee_ty = pointee.ty; |
| 93 | + if pointee_ty.is_array() || pointee_ty.is_slice() || pointee_ty.is_str() { |
| 94 | + pointee_ty = pointee_ty.sequence_element_type(self.tcx); |
| 95 | + } |
| 96 | + |
| 97 | + if !pointee_ty.is_sized(self.tcx, self.tcx.param_env_reveal_all_normalized(self.def_id)) { |
| 98 | + debug!("Unsafe pointer, but unsized: {:?}", pointer_ty); |
| 99 | + return; |
| 100 | + } |
| 101 | + |
| 102 | + if [self.tcx.types.bool, self.tcx.types.i8, self.tcx.types.u8, self.tcx.types.str_] |
| 103 | + .contains(&pointee_ty) |
| 104 | + { |
| 105 | + debug!("Trivially aligned pointee type: {:?}", pointer_ty); |
| 106 | + return; |
| 107 | + } |
| 108 | + |
| 109 | + self.pointers.push((pointer, pointee_ty)) |
| 110 | + } |
| 111 | +} |
| 112 | + |
| 113 | +fn split_block( |
| 114 | + basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'_>>, |
| 115 | + location: Location, |
| 116 | +) -> BasicBlock { |
| 117 | + let block_data = &mut basic_blocks[location.block]; |
| 118 | + |
| 119 | + // Drain every statement after this one and move the current terminator to a new basic block |
| 120 | + let new_block = BasicBlockData { |
| 121 | + statements: block_data.statements.split_off(location.statement_index), |
| 122 | + terminator: block_data.terminator.take(), |
| 123 | + is_cleanup: block_data.is_cleanup, |
| 124 | + }; |
| 125 | + |
| 126 | + basic_blocks.push(new_block) |
| 127 | +} |
| 128 | + |
| 129 | +fn insert_alignment_check<'tcx>( |
| 130 | + tcx: TyCtxt<'tcx>, |
| 131 | + local_decls: &mut LocalDecls<'tcx>, |
| 132 | + block_data: &mut BasicBlockData<'tcx>, |
| 133 | + pointer: Place<'tcx>, |
| 134 | + pointee_ty: Ty<'tcx>, |
| 135 | + source_info: SourceInfo, |
| 136 | + new_block: BasicBlock, |
| 137 | +) { |
| 138 | + // Cast the pointer to a *const () |
| 139 | + let const_raw_ptr = tcx.mk_ptr(TypeAndMut { ty: tcx.types.unit, mutbl: Mutability::Not }); |
| 140 | + let rvalue = Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(pointer), const_raw_ptr); |
| 141 | + let thin_ptr = local_decls.push(LocalDecl::with_source_info(const_raw_ptr, source_info)).into(); |
| 142 | + block_data |
| 143 | + .statements |
| 144 | + .push(Statement { source_info, kind: StatementKind::Assign(Box::new((thin_ptr, rvalue))) }); |
| 145 | + |
| 146 | + // Transmute the pointer to a usize (equivalent to `ptr.addr()`) |
| 147 | + let rvalue = Rvalue::Cast(CastKind::Transmute, Operand::Copy(thin_ptr), tcx.types.usize); |
| 148 | + let addr = local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); |
| 149 | + block_data |
| 150 | + .statements |
| 151 | + .push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) }); |
| 152 | + |
| 153 | + // Get the alignment of the pointee |
| 154 | + let alignment = |
| 155 | + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); |
| 156 | + let rvalue = Rvalue::NullaryOp(NullOp::AlignOf, pointee_ty); |
| 157 | + block_data.statements.push(Statement { |
| 158 | + source_info, |
| 159 | + kind: StatementKind::Assign(Box::new((alignment, rvalue))), |
| 160 | + }); |
| 161 | + |
| 162 | + // Subtract 1 from the alignment to get the alignment mask |
| 163 | + let alignment_mask = |
| 164 | + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); |
| 165 | + let one = Operand::Constant(Box::new(Constant { |
| 166 | + span: source_info.span, |
| 167 | + user_ty: None, |
| 168 | + literal: ConstantKind::Val( |
| 169 | + ConstValue::Scalar(Scalar::from_target_usize(1, &tcx)), |
| 170 | + tcx.types.usize, |
| 171 | + ), |
| 172 | + })); |
| 173 | + block_data.statements.push(Statement { |
| 174 | + source_info, |
| 175 | + kind: StatementKind::Assign(Box::new(( |
| 176 | + alignment_mask, |
| 177 | + Rvalue::BinaryOp(BinOp::Sub, Box::new((Operand::Copy(alignment), one))), |
| 178 | + ))), |
| 179 | + }); |
| 180 | + |
| 181 | + // BitAnd the alignment mask with the pointer |
| 182 | + let alignment_bits = |
| 183 | + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); |
| 184 | + block_data.statements.push(Statement { |
| 185 | + source_info, |
| 186 | + kind: StatementKind::Assign(Box::new(( |
| 187 | + alignment_bits, |
| 188 | + Rvalue::BinaryOp( |
| 189 | + BinOp::BitAnd, |
| 190 | + Box::new((Operand::Copy(addr), Operand::Copy(alignment_mask))), |
| 191 | + ), |
| 192 | + ))), |
| 193 | + }); |
| 194 | + |
| 195 | + // Check if the alignment bits are all zero |
| 196 | + let is_ok = local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into(); |
| 197 | + let zero = Operand::Constant(Box::new(Constant { |
| 198 | + span: source_info.span, |
| 199 | + user_ty: None, |
| 200 | + literal: ConstantKind::Val( |
| 201 | + ConstValue::Scalar(Scalar::from_target_usize(0, &tcx)), |
| 202 | + tcx.types.usize, |
| 203 | + ), |
| 204 | + })); |
| 205 | + block_data.statements.push(Statement { |
| 206 | + source_info, |
| 207 | + kind: StatementKind::Assign(Box::new(( |
| 208 | + is_ok, |
| 209 | + Rvalue::BinaryOp(BinOp::Eq, Box::new((Operand::Copy(alignment_bits), zero.clone()))), |
| 210 | + ))), |
| 211 | + }); |
| 212 | + |
| 213 | + // Set this block's terminator to our assert, continuing to new_block if we pass |
| 214 | + block_data.terminator = Some(Terminator { |
| 215 | + source_info, |
| 216 | + kind: TerminatorKind::Assert { |
| 217 | + cond: Operand::Copy(is_ok), |
| 218 | + expected: true, |
| 219 | + target: new_block, |
| 220 | + msg: AssertKind::MisalignedPointerDereference { |
| 221 | + required: Operand::Copy(alignment), |
| 222 | + found: Operand::Copy(addr), |
| 223 | + }, |
| 224 | + cleanup: None, |
| 225 | + }, |
| 226 | + }); |
| 227 | +} |
0 commit comments