1pub mod inter_visitor;
2pub mod intra_visitor;
3pub mod order;
4pub mod ownership;
5
6use rustc_middle::mir::{Body, Terminator};
7use rustc_middle::ty::{InstanceKind::Item, TyCtxt};
8use rustc_span::def_id::DefId;
9
10use super::{rCanary, IcxMut, IcxSliceMut, Rcx, RcxMut};
11use crate::analysis::core::heap_analysis::HeapInfo;
12use crate::analysis::core::heap_analysis::{default::TyWithIndex, HAResult};
13use ownership::{IntraVar, Taint};
14
15use std::{
16 collections::{HashMap, HashSet},
17 env,
18 fmt::{Debug, Formatter},
19};
20
21pub type MirGraph = HashMap<DefId, Graph>;
22pub type ToPo = Vec<usize>;
23pub type Edges = Vec<Vec<usize>>;
24
25#[derive(Debug, Clone)]
26pub struct Graph {
27 e: Edges,
28 pre: Edges,
29 topo: ToPo,
30}
31
32impl Default for Graph {
33 fn default() -> Self {
34 Self {
35 e: Vec::default(),
36 pre: Vec::default(),
37 topo: Vec::default(),
38 }
39 }
40}
41
42impl Graph {
43 pub fn new(len: usize) -> Self {
44 Graph {
45 e: vec![Vec::new(); len],
46 pre: vec![Vec::new(); len],
47 topo: Vec::new(),
48 }
49 }
50
51 pub fn get_edges(&self) -> &Edges {
52 &self.e
53 }
54
55 pub fn get_edges_mut(&mut self) -> &mut Edges {
56 &mut self.e
57 }
58
59 pub fn get_pre(&self) -> &Edges {
60 &self.pre
61 }
62
63 pub fn get_pre_mut(&mut self) -> &mut Edges {
64 &mut self.pre
65 }
66
67 pub fn get_topo(&self) -> &ToPo {
68 &self.topo
69 }
70
71 pub fn get_topo_mut(&mut self) -> &mut ToPo {
72 &mut self.topo
73 }
74}
75
76pub struct FlowAnalysis<'tcx, 'a> {
77 rcx: &'a mut rCanary<'tcx>,
78 fn_set: HashSet<DefId>,
79}
80
81impl<'tcx, 'a> FlowAnalysis<'tcx, 'a> {
82 pub fn new(rcx: &'a mut rCanary<'tcx>) -> Self {
83 Self {
84 rcx,
85 fn_set: HashSet::new(),
86 }
87 }
88
89 pub fn fn_set(&self) -> &HashSet<DefId> {
90 &self.fn_set
91 }
92
93 pub fn fn_set_mut(&mut self) -> &mut HashSet<DefId> {
94 &mut self.fn_set
95 }
96
97 pub fn mir_graph(&self) -> &MirGraph {
98 self.rcx().mir_graph()
99 }
100
101 pub fn mir_graph_mut(&mut self) -> &mut MirGraph {
102 self.rcx_mut().mir_graph_mut()
103 }
104
105 pub fn start(&mut self) {
106 self.order();
109 self.intra_run();
112 }
113}
114
115impl<'tcx, 'o, 'a> RcxMut<'tcx, 'o, 'a> for FlowAnalysis<'tcx, 'a> {
116 #[inline(always)]
117 fn rcx(&'o self) -> &'o rCanary<'tcx> {
118 self.rcx
119 }
120
121 #[inline(always)]
122 fn rcx_mut(&'o mut self) -> &'o mut rCanary<'tcx> {
123 &mut self.rcx
124 }
125
126 #[inline(always)]
127 fn tcx(&'o self) -> TyCtxt<'tcx> {
128 self.rcx().tcx()
129 }
130}
131
132#[derive(Clone, Debug)]
133pub struct NodeOrder<'tcx> {
134 body: &'tcx Body<'tcx>,
135 graph: Graph,
136}
137
138impl<'tcx> NodeOrder<'tcx> {
139 pub fn new(body: &'tcx Body<'tcx>) -> Self {
140 let len = body.basic_blocks.len();
141 Self {
142 body,
143 graph: Graph::new(len),
144 }
145 }
146
147 #[inline(always)]
148 pub fn body(&self) -> &'tcx Body<'tcx> {
149 self.body
150 }
151
152 #[inline(always)]
153 pub fn graph(&self) -> &Graph {
154 &self.graph
155 }
156
157 #[inline(always)]
158 pub fn graph_mut(&mut self) -> &mut Graph {
159 &mut self.graph
160 }
161}
162
163struct IntraFlowAnalysis<'tcx, 'ctx, 'a> {
164 pub rcx: &'a rCanary<'tcx>,
165 icx: IntraFlowContext<'tcx, 'ctx>,
166 icx_slice: IcxSliceFroBlock<'tcx, 'ctx>,
167 pub def_id: DefId,
168 pub body: &'a Body<'tcx>,
169 pub graph: &'a Graph,
170 taint_flag: bool,
171 taint_source: Vec<Terminator<'tcx>>,
172}
173
174impl<'tcx, 'ctx, 'a> IntraFlowAnalysis<'tcx, 'ctx, 'a> {
175 pub fn new(
176 rcx: &'a rCanary<'tcx>,
177 def_id: DefId,
178 ) -> Self {
180 let body = rcx.tcx.instance_mir(Item(def_id));
181 let v_len = body.local_decls.len();
182 let b_len = body.basic_blocks.len();
183 let graph = rcx.mir_graph().get(&def_id).unwrap();
184
185 Self {
186 rcx,
187 icx: IntraFlowContext::new(b_len, v_len),
188 icx_slice: IcxSliceFroBlock::new_for_block_0(v_len),
189 def_id,
190 body,
191 graph,
192 taint_flag: false,
193 taint_source: Vec::default(),
194 }
195 }
196
197 pub fn owner(&self) -> &HAResult {
198 self.rcx.adt_owner()
199 }
200
201 pub fn add_taint(&mut self, terminator: Terminator<'tcx>) {
202 self.taint_source.push(terminator);
203 }
204}
205
206impl<'tcx, 'ctx, 'o, 'a> Rcx<'tcx, 'o, 'a> for IntraFlowAnalysis<'tcx, 'ctx, 'a> {
207 #[inline(always)]
208 fn rcx(&'o self) -> &'a rCanary<'tcx> {
209 self.rcx
210 }
211
212 #[inline(always)]
213 fn tcx(&'o self) -> TyCtxt<'tcx> {
214 self.rcx.tcx()
215 }
216}
217
218impl<'tcx, 'ctx, 'o, 'a> IcxMut<'tcx, 'ctx, 'o> for IntraFlowAnalysis<'tcx, 'ctx, 'a> {
219 #[inline(always)]
220 fn icx(&'o self) -> &'o IntraFlowContext<'tcx, 'ctx> {
221 &self.icx
222 }
223
224 #[inline(always)]
225 fn icx_mut(&'o mut self) -> &'o mut IntraFlowContext<'tcx, 'ctx> {
226 &mut self.icx
227 }
228}
229
230impl<'tcx, 'ctx, 'o, 'a> IcxSliceMut<'tcx, 'ctx, 'o> for IntraFlowAnalysis<'tcx, 'ctx, 'a> {
231 #[inline(always)]
232 fn icx_slice(&'o self) -> &'o IcxSliceFroBlock<'tcx, 'ctx> {
233 &self.icx_slice
234 }
235
236 #[inline(always)]
237 fn icx_slice_mut(&'o mut self) -> &'o mut IcxSliceFroBlock<'tcx, 'ctx> {
238 &mut self.icx_slice
239 }
240}
241
242#[derive(Debug, Clone)]
243pub struct IntraFlowContext<'tcx, 'ctx> {
244 taint: IOPairForGraph<Taint<'tcx>>,
245 var: IOPairForGraph<IntraVar<'ctx>>,
246 len: IOPairForGraph<usize>,
247 ty: IOPairForGraph<TyWithIndex<'tcx>>,
250 layout: IOPairForGraph<Vec<HeapInfo>>,
251}
252
253impl<'tcx, 'ctx, 'icx> IntraFlowContext<'tcx, 'ctx> {
254 pub fn new(b_len: usize, v_len: usize) -> Self {
255 Self {
256 taint: IOPairForGraph::new(b_len, v_len),
257 var: IOPairForGraph::new(b_len, v_len),
258 len: IOPairForGraph::new(b_len, v_len),
259 ty: IOPairForGraph::new(b_len, v_len),
260 layout: IOPairForGraph::new(b_len, v_len),
261 }
262 }
263
264 pub fn taint(&self) -> &IOPairForGraph<Taint<'tcx>> {
265 &self.taint
266 }
267
268 pub fn taint_mut(&mut self) -> &mut IOPairForGraph<Taint<'tcx>> {
269 &mut self.taint
270 }
271
272 pub fn var(&self) -> &IOPairForGraph<IntraVar<'ctx>> {
273 &self.var
274 }
275
276 pub fn var_mut(&mut self) -> &mut IOPairForGraph<IntraVar<'ctx>> {
277 &mut self.var
278 }
279
280 pub fn len(&self) -> &IOPairForGraph<usize> {
281 &self.len
282 }
283
284 pub fn len_mut(&mut self) -> &mut IOPairForGraph<usize> {
285 &mut self.len
286 }
287
288 pub fn ty(&self) -> &IOPairForGraph<TyWithIndex<'tcx>> {
289 &self.ty
290 }
291
292 pub fn ty_mut(&mut self) -> &mut IOPairForGraph<TyWithIndex<'tcx>> {
293 &mut self.ty
294 }
295
296 pub fn layout(&self) -> &IOPairForGraph<Vec<HeapInfo>> {
297 &self.layout
298 }
299
300 pub fn layout_mut(&mut self) -> &mut IOPairForGraph<Vec<HeapInfo>> {
301 &mut self.layout
302 }
303
304 pub fn derive_from_pre_node(&mut self, from: usize, to: usize) {
305 *self.taint_mut().get_g_mut()[to].get_i_mut() =
307 self.taint_mut().get_g_mut()[from].get_o_mut().clone();
308
309 *self.var_mut().get_g_mut()[to].get_i_mut() =
311 self.var_mut().get_g_mut()[from].get_o_mut().clone();
312
313 *self.len_mut().get_g_mut()[to].get_i_mut() =
315 self.len_mut().get_g_mut()[from].get_o_mut().clone();
316
317 *self.ty_mut().get_g_mut()[to].get_i_mut() =
319 self.ty_mut().get_g_mut()[from].get_o_mut().clone();
320
321 *self.layout_mut().get_g_mut()[to].get_i_mut() =
323 self.layout_mut().get_g_mut()[from].get_o_mut().clone();
324 }
325
326 pub fn derive_from_icx_slice(&mut self, from: IcxSliceFroBlock<'tcx, 'ctx>, to: usize) {
327 *self.taint_mut().get_g_mut()[to].get_o_mut() = from.taint;
328
329 *self.var_mut().get_g_mut()[to].get_o_mut() = from.var;
330
331 *self.len_mut().get_g_mut()[to].get_o_mut() = from.len;
332
333 *self.ty_mut().get_g_mut()[to].get_o_mut() = from.ty;
334
335 *self.layout_mut().get_g_mut()[to].get_o_mut() = from.layout;
336 }
337}
338
339#[derive(Debug, Clone, Default)]
340pub struct InOutPair<T: Debug + Clone + Default> {
341 i: Vec<T>,
342 o: Vec<T>,
343}
344
345impl<T> InOutPair<T>
346where
347 T: Debug + Clone + Default,
348{
349 pub fn new(len: usize) -> Self {
350 Self {
351 i: vec![T::default(); len],
352 o: vec![T::default(); len],
353 }
354 }
355
356 pub fn get_i(&self) -> &Vec<T> {
357 &self.i
358 }
359
360 pub fn get_o(&self) -> &Vec<T> {
361 &self.o
362 }
363
364 pub fn get_i_mut(&mut self) -> &mut Vec<T> {
365 &mut self.i
366 }
367
368 pub fn get_o_mut(&mut self) -> &mut Vec<T> {
369 &mut self.o
370 }
371
372 pub fn len(&self) -> usize {
373 self.i.len()
374 }
375}
376
377#[derive(Debug, Clone, Default)]
378pub struct IOPairForGraph<T: Debug + Clone + Default> {
379 pair_graph: Vec<InOutPair<T>>,
380}
381
382impl<T> IOPairForGraph<T>
383where
384 T: Debug + Clone + Default,
385{
386 pub fn new(b_len: usize, v_len: usize) -> Self {
387 Self {
388 pair_graph: vec![InOutPair::new(v_len); b_len],
389 }
390 }
391
392 pub fn get_g(&self) -> &Vec<InOutPair<T>> {
393 &self.pair_graph
394 }
395
396 pub fn get_g_mut(&mut self) -> &mut Vec<InOutPair<T>> {
397 &mut self.pair_graph
398 }
399}
400
401#[derive(Clone, Default)]
402pub struct IcxSliceFroBlock<'tcx, 'ctx> {
403 taint: Vec<Taint<'tcx>>,
404 var: Vec<IntraVar<'ctx>>,
405 len: Vec<usize>,
406 ty: Vec<TyWithIndex<'tcx>>,
409 layout: Vec<Vec<HeapInfo>>,
410}
411
412impl<'tcx, 'ctx> IcxSliceFroBlock<'tcx, 'ctx> {
413 pub fn new_in(icx: &mut IntraFlowContext<'tcx, 'ctx>, idx: usize) -> Self {
414 Self {
415 taint: icx.taint_mut().get_g_mut()[idx].get_i_mut().clone(),
416 var: icx.var_mut().get_g_mut()[idx].get_i_mut().clone(),
417 len: icx.len_mut().get_g_mut()[idx].get_i_mut().clone(),
418 ty: icx.ty_mut().get_g_mut()[idx].get_i_mut().clone(),
419 layout: icx.layout_mut().get_g_mut()[idx].get_i_mut().clone(),
420 }
421 }
422
423 pub fn new_out(icx: &mut IntraFlowContext<'tcx, 'ctx>, idx: usize) -> Self {
424 Self {
425 taint: icx.taint_mut().get_g_mut()[idx].get_o_mut().clone(),
426 var: icx.var_mut().get_g_mut()[idx].get_o_mut().clone(),
427 len: icx.len_mut().get_g_mut()[idx].get_o_mut().clone(),
428 ty: icx.ty_mut().get_g_mut()[idx].get_o_mut().clone(),
429 layout: icx.layout_mut().get_g_mut()[idx].get_o_mut().clone(),
430 }
431 }
432
433 pub fn new_for_block_0(len: usize) -> Self {
434 Self {
435 taint: vec![Taint::default(); len],
436 var: vec![IntraVar::default(); len],
437 len: vec![0; len],
438 ty: vec![TyWithIndex::default(); len],
439 layout: vec![Vec::new(); len],
440 }
441 }
442
443 pub fn taint(&self) -> &Vec<Taint<'tcx>> {
444 &self.taint
445 }
446
447 pub fn taint_mut(&mut self) -> &mut Vec<Taint<'tcx>> {
448 &mut self.taint
449 }
450
451 pub fn var(&self) -> &Vec<IntraVar<'ctx>> {
452 &self.var
453 }
454
455 pub fn var_mut(&mut self) -> &mut Vec<IntraVar<'ctx>> {
456 &mut self.var
457 }
458
459 pub fn len(&self) -> &Vec<usize> {
460 &self.len
461 }
462
463 pub fn len_mut(&mut self) -> &mut Vec<usize> {
464 &mut self.len
465 }
466
467 pub fn ty(&self) -> &Vec<TyWithIndex<'tcx>> {
468 &self.ty
469 }
470
471 pub fn ty_mut(&mut self) -> &mut Vec<TyWithIndex<'tcx>> {
472 &mut self.ty
473 }
474
475 pub fn layout(&self) -> &Vec<Vec<HeapInfo>> {
476 &self.layout
477 }
478
479 pub fn layout_mut(&mut self) -> &mut Vec<Vec<HeapInfo>> {
480 &mut self.layout
481 }
482
483 pub fn taint_merge(&mut self, another: &IcxSliceFroBlock<'tcx, 'ctx>, u: usize) {
484 if another.taint()[u].is_untainted() {
485 return;
486 }
487
488 if self.taint()[u].is_untainted() {
489 self.taint_mut()[u] = another.taint()[u].clone();
490 } else {
491 for elem in another.taint()[u].set().clone() {
492 self.taint_mut()[u].insert(elem);
493 }
494 }
495 }
496}
497
498impl<'tcx, 'ctx> Debug for IcxSliceFroBlock<'tcx, 'ctx> {
499 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
500 write!(
501 f,
502 "IcxSliceForBlock\n {:?}\n {:?}\n {:?}\n {:?}\n {:?}",
503 self.taint(),
504 self.len(),
505 self.var(),
506 self.layout(),
507 self.ty(),
508 )
509 }
510}
511
512#[derive(Debug, Copy, Clone, Hash)]
513pub enum Z3GoalDisplay {
514 Verbose,
515 Disabled,
516}
517
518pub fn is_z3_goal_verbose() -> bool {
519 match env::var_os("Z3") {
520 Some(_) => true,
521 _ => false,
522 }
523}
524
525#[derive(Debug, Copy, Clone, Hash)]
526pub enum IcxSliceDisplay {
527 Verbose,
528 Disabled,
529}
530
531pub fn is_icx_slice_verbose() -> bool {
532 match env::var_os("ICX_SLICE") {
533 Some(_) => true,
534 _ => false,
535 }
536}