From 57e1f7f4e601d697e68f07a9b0039eebf59d761e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Wed, 28 Apr 2021 17:39:03 +0200
Subject: [PATCH] Adds reschedule_instructions()

---
 src/jit.rs             | 18 +++++++++++++++
 src/static_analysis.rs | 52 +++++++++++++++++++++++++++---------------
 2 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/src/jit.rs b/src/jit.rs
index 7b0d3c429..922598899 100644
--- a/src/jit.rs
+++ b/src/jit.rs
@@ -107,6 +107,24 @@ fn _scramble_instructions_in_basic_block<E: UserDefinedError, I: InstructionMete
     schedule
 }
 
+/// Reassembles the instructions in an executable according to the given schedule.
+fn _reschedule_instructions<E: UserDefinedError, I: 'static + InstructionMeter>(
+    static_analysis: &Analysis<E, I>,
+    schedule: &[usize],
+) -> Box<dyn Executable<E, I>> {
+    let (_program_vm_addr, input_program) = static_analysis.executable.get_text_bytes().unwrap();
+    let mut output_program = Vec::new();
+    let mut output_offset = 0;
+    for insn_pc in schedule {
+        let input_offset = *insn_pc * ebpf::INSN_SIZE;
+        let insn = ebpf::get_insn_unchecked(input_program, *insn_pc);
+        let length = if insn.opc == ebpf::LD_DW_IMM { 2 } else { 1 } * ebpf::INSN_SIZE;
+        output_program[output_offset..output_offset + length].clone_from_slice(&input_program[input_offset..input_offset + length]);
+        output_offset += length;
+    }
+    <dyn Executable<E, I>>::from_text_bytes(&output_program, None, *static_analysis.executable.get_config()).unwrap()
+}
+
 /// Argument for executing a eBPF JIT-compiled program
 pub struct JitProgramArgument<'a> {
     /// The MemoryMapping to be used to run the compiled code
diff --git a/src/static_analysis.rs b/src/static_analysis.rs
index a2ab5f9fd..51f5b7daa 100644
--- a/src/static_analysis.rs
+++ b/src/static_analysis.rs
@@ -424,6 +424,25 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> Analysis<'a, E, I> {
         Ok(())
     }
 
+    /// Iterates over the cfg_nodes while providing the PC range of the function they belong to.
+    pub fn iter_cfg_by_function(
+        &self,
+    ) -> impl Iterator<Item = (std::ops::Range<usize>, usize)> + '_ {
+        let mut function_iter = self.functions.iter().peekable();
+        let mut function_start = *function_iter.next().unwrap().0;
+        self.cfg_nodes.keys().map(move |cfg_node_start| {
+            if function_iter.peek().unwrap().0 == cfg_node_start {
+                function_start = *function_iter.next().unwrap().0;
+            }
+            let function_end = if let Some(next_function) = function_iter.peek() {
+                *next_function.0
+            } else {
+                self.instructions.last().unwrap().ptr + 1
+            };
+            (function_start..function_end, *cfg_node_start)
+        })
+    }
+
     /// Generates a graphviz DOT of the analyzed executable
     pub fn visualize_graphically<W: std::io::Write>(
         &self,
@@ -538,22 +557,13 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> Analysis<'a, E, I> {
             }
             writeln!(output, "  }}")?;
         }
-        let mut function_iter = self.functions.iter().peekable();
-        let mut function_start = *function_iter.next().unwrap().0;
-        for (cfg_node_start, cfg_node) in self.cfg_nodes.iter() {
-            if function_iter.peek().unwrap().0 == cfg_node_start {
-                function_start = *function_iter.next().unwrap().0;
-            }
-            let function_end = if let Some(next_function) = function_iter.peek() {
-                *next_function.0
-            } else {
-                self.instructions.last().unwrap().ptr + 1
-            };
-            if *cfg_node_start != cfg_node.dominator_parent {
+        for (function_range, cfg_node_start) in self.iter_cfg_by_function() {
+            let cfg_node = &self.cfg_nodes[&cfg_node_start];
+            if cfg_node_start != cfg_node.dominator_parent {
                 writeln!(
                     output,
                     "  lbb_{} -> lbb_{} [style=dotted; arrowhead=none];",
-                    *cfg_node_start, cfg_node.dominator_parent,
+                    cfg_node_start, cfg_node.dominator_parent,
                 )?;
             }
             let mut edges: BTreeMap<usize, usize> = cfg_node
@@ -562,7 +572,7 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> Analysis<'a, E, I> {
                 .map(|destination| (*destination, 0))
                 .collect();
             if let Some(dynamic_analysis) = dynamic_analysis {
-                if let Some(recorded_edges) = dynamic_analysis.edges.get(cfg_node_start) {
+                if let Some(recorded_edges) = dynamic_analysis.edges.get(&cfg_node_start) {
                     for (destination, recorded_counter) in recorded_edges.iter() {
                         edges
                             .entry(*destination)
@@ -578,7 +588,7 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> Analysis<'a, E, I> {
                 writeln!(
                     output,
                     "  lbb_{} -> {{{}}};",
-                    *cfg_node_start,
+                    cfg_node_start,
                     edges
                         .keys()
                         .map(|destination| format!("lbb_{}", *destination))
@@ -588,11 +598,15 @@ impl<'a, E: UserDefinedError, I: InstructionMeter> Analysis<'a, E, I> {
             } else {
                 let dynamic_analysis = dynamic_analysis.unwrap();
                 for (destination, counter) in edges {
-                    write!(output, "  lbb_{} -> ", *cfg_node_start,)?;
-                    if (function_start..function_end).contains(&destination) {
-                        write!(output, "lbb_{}", destination,)?;
+                    write!(output, "  lbb_{} -> ", cfg_node_start)?;
+                    if function_range.contains(&destination) {
+                        write!(output, "lbb_{}", destination)?;
                     } else {
-                        write!(output, "alias_{0}_lbb_{1}", function_start, destination)?;
+                        write!(
+                            output,
+                            "alias_{0}_lbb_{1}",
+                            function_range.start, destination
+                        )?;
                     }
                     writeln!(
                         output,