graph generator

libafl/src/events/mod.rs

@@ -63,7 +63,7 @@ use crate::{
 };
 
 /// Multi-machine mode
-#[cfg(feature = "multi_machine")]
+#[cfg(all(unix, feature = "std", feature = "multi_machine"))]
 pub mod multi_machine;
 
 /// Check if ctrl-c is sent with this struct

utils/multi_machine_generator/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "multi_machine_generator"
+authors = ["Romain Malmain <[email protected]>"]
+edition = "2021"
+
+[dependencies]
+petgraph = "0.6"

utils/multi_machine_generator/src/graph.rs

@@ -0,0 +1,75 @@
+use std::net::SocketAddr;
+use petgraph::{Direction, Graph};
+use petgraph::graph::NodeIndex;
+
+/// A node of the network
+#[derive(Debug, Clone)]
+pub struct MultiMachineNode {
+    addr: SocketAddr
+}
+
+/// The tree
+pub struct MultiMachineTree {
+    pub graph: Graph<MultiMachineNode, ()>
+}
+
+impl MultiMachineNode {
+    pub fn new(addr: SocketAddr) -> Self {
+        Self {
+            addr
+        }
+    }
+}
+
+impl MultiMachineTree {
+    /// Generate a multi-machine tree.
+    ///
+    ///
+    /// - machines: machines to add.
+    /// - max_children_per_parent: each parent will have at most this amount of children
+    pub fn generate(machines: &[SocketAddr], max_children_per_parent: u64) -> Self {
+        let mut graph = Graph::<MultiMachineNode, ()>::new();
+        let mut machines = Vec::from(machines);
+
+        let root = if let Some(root) = machines.pop() {
+            graph.add_node(MultiMachineNode::new(root))
+        } else {
+            return Self {
+                graph
+            };
+        };
+
+        let mut graph = Self { graph
+        };
+
+        let mut populate_idx = 0u64; // round-robin population to avoid congestion
+        let mut nodes_to_populate_now: Vec<NodeIndex> = vec![root]; // current nodes we are working on
+
+        let mut nodes_to_populate_later: Vec<NodeIndex> = Vec::new();
+
+        // place all the machines in the graph
+        while let Some(machine) = machines.pop() {
+            if graph.nb_children(nodes_to_populate_now[populate_idx as usize]) == max_children_per_parent {
+                nodes_to_populate_now = nodes_to_populate_later.drain(..).collect();
+                populate_idx = 0; // should be useless
+            }
+
+            let new_child = graph.add_child(nodes_to_populate_now[populate_idx as usize], MultiMachineNode::new(machine));
+            nodes_to_populate_later.push(new_child);
+
+            populate_idx = (populate_idx + 1) % nodes_to_populate_now.len() as u64;
+        }
+
+        graph
+    }
+
+    fn add_child(&mut self, parent: NodeIndex, child: MultiMachineNode) -> NodeIndex {
+        let child_idx = self.graph.add_node(child);
+        self.graph.add_edge(child_idx, parent, ());
+        child_idx
+    }
+
+    fn nb_children(&self, node: NodeIndex) -> u64 {
+        self.graph.neighbors_directed(node, Direction::Incoming).count() as u64
+    }
+}

utils/multi_machine_generator/src/main.rs

@@ -0,0 +1,45 @@
+//! Multi Machine Generator
+//!
+//! Generates a ready-to-run multi-machine configuration, as a balanced tree.
+//! A simple algorithm will first create such a tree, and associate IPs to them.
+//! It will finally output a set of commands to run to have each fuzzer communicating correctly with the other machines of the network.
+//!
+//! We suppose everyone is on the same network and the machines have the fuzzer ready to run on each machine.
+
+
+use std::fs;
+use std::net::SocketAddr;
+use std::str::FromStr;
+use petgraph::dot::Dot;
+use crate::graph::MultiMachineTree;
+
+pub mod graph;
+
+fn main() {
+    let machines = [
+        SocketAddr::from_str("0.0.0.1:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.2:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.3:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.4:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.5:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.6:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.7:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.8:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.9:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.10:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.11:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.12:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.13:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.14:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.15:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.16:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.17:50000").unwrap(),
+        SocketAddr::from_str("0.0.0.18:50000").unwrap(),
+    ];
+
+    let multi_machine_graph = MultiMachineTree::generate(&machines, 3);
+
+    let dot = Dot::new(&multi_machine_graph.graph);
+
+    fs::write("multi_machine.dot", format!("{:?}", dot)).unwrap();
+}