libinject/
trajectory.rs

1use crate::{drcore::log, pipe};
2use navigator::{
3    netgraph::{
4        core::{Attempt, payload_softmatch_attempt},
5        trajectory::{GraphStep, PlanningStep, Trajectory},
6    },
7    pipe_event::PipeEvent,
8};
9use spin::Mutex;
10use std::sync::LazyLock;
11
12static TRAJECTORY: LazyLock<Mutex<Option<Trajectory<PlanningStep>>>> =
13    LazyLock::new(|| Mutex::new(None));
14/// The sub-trajectory that is repetedly traversed during fuzzing with soft-rollbacks. This
15/// sub-trajectory must be recorded at runtime, as it is dependent on the fuzzing target offset.
16static ROLLBACK_SUB_TRAJECTORY: LazyLock<Mutex<Option<Trajectory<PlanningStep>>>> =
17    LazyLock::new(|| Mutex::new(None));
18
19pub fn init(config_trajectory: Trajectory<PlanningStep>) {
20    let mut t = TRAJECTORY.lock();
21    *t = Some(config_trajectory.clone());
22}
23
24#[derive(Debug)]
25pub enum TrajectoryError {
26    OffTrajectory,
27    InvalidTrajectory,
28    EndOfTrajectory,
29}
30
31/// Record the sub-trajectory that we expect to repetedely traverse during fuzzing soft-rollback.
32pub fn record_rollback_sub_trajectory() {
33    let mut rollback_sub_t = ROLLBACK_SUB_TRAJECTORY.lock();
34
35    if let None = *rollback_sub_t {
36        let current_t = TRAJECTORY
37            .lock()
38            .as_ref()
39            .expect("Trajectory must be Some on a fuzzing run, and initialised first")
40            .clone();
41
42        log(&format!(
43            "[trajectory] recorded rollback_sub_trajectory: {:?}",
44            current_t
45        ));
46
47        *rollback_sub_t = Some(current_t);
48    }
49}
50
51/// Restore TRAJECTORY to the sub-trajectory that was recorded with
52/// `record_rollback_sub_trajectory`.
53pub fn restore_rollback_sub_trajectory() {
54    let sub_trajectory = ROLLBACK_SUB_TRAJECTORY
55        .lock()
56        .as_ref()
57        .expect("record_rollback_sub_trajectory must be called first")
58        .clone();
59
60    log(&format!(
61        "[trajectory] restored trajectory to the recorded rollback_sub_trajectory: {:?}",
62        sub_trajectory
63    ));
64
65    let mut t = TRAJECTORY.lock();
66    *t = Some(sub_trajectory);
67}
68
69/// Pops a step from the front of the TRAJECTORY.
70/// Panics if:
71/// - Unable to get a lock on the TRAJECTORY.
72/// - The trajectory module has not first been initialised with `trajectory::init`.
73pub fn pop_front() -> Option<PlanningStep> {
74    let mut binding = TRAJECTORY.lock();
75    binding
76        .as_mut()
77        .and_then(Trajectory::<PlanningStep>::pop_front)
78}
79
80/// Tries to pop a step from the front of the TRAJECTORY. If the trajectory is empty, `Ok(None)` is returned.
81/// Otherwise checks that the popped `planning_step` matches the provided `Attempt`, and then returns a new
82/// `PlanningStep` (as an `Ok(Some(planning_step))`), with the `Attempt` in the popped step replaced with the **provided** `Attempt`.
83/// If the `Attempt`s do not match, `Err(TrajectoryError::OffTrajectory)` is returned and a
84/// `TrajectoryDeviation` is sent up the pipe to the navigator.
85pub fn next(observed_attempt: &Attempt) -> Result<Option<PlanningStep>, TrajectoryError> {
86    match pop_front() {
87        None => {
88            log(&format!("[trajectory] Empty trajectory"));
89            Ok(None)
90        }
91        Some(planning_step) => {
92            let planned_attempt = match planning_step.clone() {
93                PlanningStep::Fuzz(recv_attempt) => Attempt::Recv(recv_attempt),
94                PlanningStep::Graph(graph_step) => graph_step.into_tuple().0,
95            };
96
97            if payload_softmatch_attempt(&planned_attempt, observed_attempt) {
98                Ok(Some(match planning_step {
99                    PlanningStep::Graph(graph_step) => PlanningStep::Graph(
100                        GraphStep::new(observed_attempt.clone(), graph_step.into_tuple().1).expect(
101                            "payload_softmatch_attempt will only match if the attempt types match",
102                        ),
103                    ),
104                    PlanningStep::Fuzz(_) => {
105                        let Attempt::Recv(recv_attempt) = observed_attempt else {
106                            unreachable!(
107                                "observed attempt has matched the recv_attempt in the planned step"
108                            )
109                        };
110                        PlanningStep::Fuzz(recv_attempt.clone())
111                    }
112                }))
113            } else {
114                log(&format!(
115                    "[trajectory] Off trajectory, found {:?}, expected {:?}",
116                    planned_attempt, observed_attempt
117                ));
118                pipe::send(&PipeEvent::TrajectoryDeviation);
119                Err(TrajectoryError::OffTrajectory)
120            }
121        }
122    }
123}
libinject/trajectory.rs

libinject/
trajectory.rs
