crackers/synthesis/
mod.rs

1use jingle::modeling::ModeledInstruction;
2use std::cmp::Ordering;
3use std::sync::Arc;
4use tracing::{Level, event, instrument};
5use z3::Context;
6
7use crate::error::CrackersError;
8use crate::error::CrackersError::EmptySpecification;
9use crate::gadget::candidates::{CandidateBuilder, Candidates};
10use crate::gadget::library::GadgetLibrary;
11use crate::reference_program::ReferenceProgram;
12use crate::synthesis::assignment_model::builder::AssignmentModelBuilder;
13use crate::synthesis::builder::{
14    StateConstraintGenerator, SynthesisParams, SynthesisSelectionStrategy,
15    TransitionConstraintGenerator,
16};
17use crate::synthesis::pcode_theory::builder::PcodeTheoryBuilder;
18use crate::synthesis::pcode_theory::theory_worker::TheoryWorker;
19use crate::synthesis::selection_strategy::AssignmentResult::{Failure, Success};
20use crate::synthesis::selection_strategy::OuterProblem::{OptimizeProb, SatProb};
21use crate::synthesis::selection_strategy::optimization_problem::OptimizationProblem;
22use crate::synthesis::selection_strategy::sat_problem::SatProblem;
23use crate::synthesis::selection_strategy::{OuterProblem, SelectionFailure, SelectionStrategy};
24use crate::synthesis::slot_assignments::SlotAssignments;
25
26pub mod assignment_model;
27pub mod builder;
28mod combined;
29pub(crate) mod partition_iterator;
30pub mod pcode_theory;
31pub mod selection_strategy;
32pub mod slot_assignments;
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
35pub struct Decision {
36    pub index: usize,
37    pub choice: usize,
38}
39
40impl PartialOrd for Decision {
41    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
42        self.index.partial_cmp(&other.index)
43    }
44}
45
46#[derive(Debug)]
47pub enum DecisionResult {
48    AssignmentFound(AssignmentModelBuilder),
49    Unsat(SelectionFailure),
50}
51
52pub struct AssignmentSynthesis {
53    outer_problem: OuterProblem,
54    library: Arc<GadgetLibrary>,
55    candidates: Candidates,
56    pointer_invariants: Vec<Arc<TransitionConstraintGenerator>>,
57    preconditions: Vec<Arc<StateConstraintGenerator>>,
58    postconditions: Vec<Arc<StateConstraintGenerator>>,
59    candidates_per_slot: usize,
60    instructions: ReferenceProgram,
61    parallel: usize,
62}
63
64impl AssignmentSynthesis {
65    pub fn new(builder: &SynthesisParams) -> Result<Self, CrackersError> {
66        let instrs = &builder.reference_program;
67        if instrs.is_empty() {
68            return Err(EmptySpecification);
69        }
70        let arch_info = &builder.gadget_library.arch_info();
71        let modeled_instrs: Vec<ModeledInstruction> = instrs
72            .steps()
73            .iter()
74            .map(|i| i.model(arch_info).unwrap())
75            .collect();
76
77        let candidates = CandidateBuilder::default()
78            .with_random_sample_size(builder.candidates_per_slot)
79            .build(builder.gadget_library.get_random_candidates_for_trace(
80                arch_info,
81                modeled_instrs.as_slice(),
82                builder.seed,
83            ))?;
84        let outer_problem = match builder.selection_strategy {
85            SynthesisSelectionStrategy::SatStrategy => {
86                SatProb(SatProblem::initialize(&candidates.candidates))
87            }
88            SynthesisSelectionStrategy::OptimizeStrategy => {
89                OptimizeProb(OptimizationProblem::initialize(&candidates.candidates))
90            }
91        };
92        Ok(AssignmentSynthesis {
93            outer_problem,
94            candidates,
95            library: builder.gadget_library.clone(),
96            pointer_invariants: builder.pointer_invariants.clone(),
97            preconditions: builder.preconditions.clone(),
98            postconditions: builder.postconditions.clone(),
99            candidates_per_slot: builder.candidates_per_slot,
100            instructions: builder.reference_program.clone(),
101            parallel: builder.parallel,
102        })
103    }
104
105    fn make_model_builder(&self, slot_assignments: SlotAssignments) -> AssignmentModelBuilder {
106        AssignmentModelBuilder {
107            templates: self.instructions.clone(),
108            gadgets: slot_assignments.interpret_from_library(&self.candidates),
109            preconditions: self.preconditions.clone(),
110            postconditions: self.postconditions.clone(),
111            pointer_invariants: self.pointer_invariants.clone(),
112            arch_info: self.library.arch_info(),
113        }
114    }
115
116    fn make_pcode_theory_builder(&self) -> PcodeTheoryBuilder<'_> {
117        PcodeTheoryBuilder::new(self.candidates.clone(), &self.library)
118            .with_pointer_invariants(&self.pointer_invariants)
119            .with_preconditions(&self.preconditions)
120            .with_postconditions(&self.postconditions)
121            .with_max_candidates(self.candidates_per_slot)
122            .with_templates(self.instructions.clone())
123    }
124
125    pub fn decide_single_threaded(&mut self) -> Result<DecisionResult, CrackersError> {
126        let theory_builder = self.make_pcode_theory_builder();
127        let theory = theory_builder.build()?;
128        loop {
129            let assignment = self.outer_problem.get_assignments()?;
130            match assignment {
131                Success(a) => {
132                    let theory_result = theory.check_assignment(&a)?;
133                    match theory_result {
134                        None => {
135                            // success
136                            return Ok(DecisionResult::AssignmentFound(self.make_model_builder(a)));
137                        }
138                        Some(conflict) => {
139                            self.outer_problem.add_theory_clauses(&conflict);
140                        }
141                    }
142                }
143                Failure(d) => return Ok(DecisionResult::Unsat(d)),
144            }
145        }
146    }
147    #[instrument(skip_all)]
148    pub fn decide(&mut self) -> Result<DecisionResult, CrackersError> {
149        let mut req_channels = vec![];
150        let mut kill_senders = vec![];
151        let library = self.library.clone();
152        let theory_builder = PcodeTheoryBuilder::new(self.candidates.clone(), &library)
153            .with_pointer_invariants(&self.pointer_invariants)
154            .with_preconditions(&self.preconditions)
155            .with_postconditions(&self.postconditions)
156            .with_max_candidates(self.candidates_per_slot)
157            .with_templates(self.instructions.clone());
158        let (resp_sender, resp_receiver) = std::sync::mpsc::channel();
159        std::thread::scope(|s| {
160            for idx in 0..self.parallel {
161                let t = theory_builder.clone();
162                let r = resp_sender.clone();
163                let (req_sender, req_receiver) = std::sync::mpsc::channel();
164                let (kill_sender, kill_receiver) = std::sync::mpsc::channel();
165                kill_senders.push(kill_sender);
166                req_channels.push(req_sender);
167                s.spawn(move || {
168                    let z3 = Context::thread_local();
169                    std::thread::scope(|inner| {
170                        let handle = z3.handle();
171                        inner.spawn(move || {
172                            for _ in kill_receiver {
173                                handle.interrupt();
174                            }
175                        });
176                        let worker = TheoryWorker::new(idx, r, req_receiver, t).unwrap();
177                        event!(Level::TRACE, "Created worker {}", idx);
178                        worker.run();
179                        drop(worker);
180                    });
181                });
182            }
183            drop(resp_sender);
184            for (i, x) in req_channels.iter().enumerate() {
185                event!(
186                    Level::TRACE,
187                    "Asking outer procedure for initial assignments"
188                );
189                if let Ok(assignment) = self.outer_problem.get_assignments() {
190                    match assignment {
191                        Success(assignment) => {
192                            event!(Level::TRACE, "Sending {:?} to worker {}", &assignment, i);
193                            x.send(assignment).unwrap();
194                        }
195                        Failure(a) => {
196                            req_channels.clear();
197                            for x in &kill_senders {
198                                x.send(()).unwrap();
199                            }
200                            kill_senders.clear();
201                            return Ok(DecisionResult::Unsat(a));
202                        }
203                    }
204                }
205            }
206            event!(Level::TRACE, "Done sending initial jobs");
207
208            for response in resp_receiver {
209                event!(
210                    Level::TRACE,
211                    "Received response from worker {}",
212                    response.idx
213                );
214
215                match response.theory_result {
216                    Ok(r) => {
217                        match r {
218                            None => {
219                                event!(
220                                    Level::INFO,
221                                    "Theory returned SAT for {:?}!",
222                                    response.assignment
223                                );
224                                req_channels.clear();
225                                for x in &kill_senders {
226                                    x.send(()).unwrap();
227                                }
228                                kill_senders.clear();
229                                return Ok(DecisionResult::AssignmentFound(
230                                    self.make_model_builder(response.assignment),
231                                ));
232                            }
233                            Some(c) => {
234                                event!(
235                                    Level::TRACE,
236                                    "Worker {} found conflicts: {}",
237                                    response.idx,
238                                    response.assignment.display_conflict(&c)
239                                );
240                                self.outer_problem.add_theory_clauses(&c);
241                                let new_assignment = self.outer_problem.get_assignments()?;
242                                match new_assignment {
243                                    Failure(a) => {
244                                        // drop the senders
245                                        req_channels.clear();
246                                        for x in &kill_senders {
247                                            x.send(()).unwrap();
248                                        }
249                                        kill_senders.clear();
250                                        return Ok(DecisionResult::Unsat(a));
251                                    }
252                                    Success(a) => {
253                                        req_channels[response.idx].send(a).unwrap();
254                                    }
255                                }
256                            }
257                        }
258                    }
259                    Err(e) => {
260                        event!(
261                            Level::ERROR,
262                            "Worker {} returned error: {}",
263                            response.idx,
264                            e
265                        );
266                        std::process::exit(-1);
267                    }
268                }
269            }
270            event!(
271                Level::ERROR,
272                "Outer SAT returned UNSAT! No solution found! :("
273            );
274            unreachable!()
275        })
276    }
277}