Add repair message support to dos tool (#9090)

Cargo.lock

@@ -4023,6 +4023,7 @@ dependencies = [
  "solana-logger 1.2.0",
  "solana-net-utils 1.2.0",
  "solana-runtime 1.2.0",
+ "solana-sdk 1.2.0",
 ]
 
 [[package]]

core/src/serve_repair.rs

@@ -59,7 +59,7 @@ pub struct ServeRepairStats {
 
 /// Window protocol messages
 #[derive(Serialize, Deserialize, Debug)]
-enum RepairProtocol {
+pub enum RepairProtocol {
     WindowIndex(ContactInfo, u64, u64),
     HighestWindowIndex(ContactInfo, u64, u64),
     Orphan(ContactInfo, u64),

dos/Cargo.toml

@@ -17,3 +17,4 @@ solana-core = { path = "../core", version = "1.2.0" }
 solana-logger = { path = "../logger", version = "1.2.0" }
 solana-net-utils = { path = "../net-utils", version = "1.2.0" }
 solana-runtime = { path = "../runtime", version = "1.2.0" }
+solana-sdk = { path = "../sdk", version = "1.2.0" }

dos/src/main.rs

@@ -1,11 +1,95 @@
 use clap::{value_t, value_t_or_exit, App, Arg};
 use log::*;
 use rand::{thread_rng, Rng};
-use solana_core::gossip_service::discover_cluster;
+use solana_core::contact_info::ContactInfo;
+use solana_core::gossip_service::discover;
+use solana_core::serve_repair::RepairProtocol;
+use solana_sdk::pubkey::Pubkey;
 use std::net::{SocketAddr, UdpSocket};
 use std::process::exit;
 use std::time::Instant;
 
+fn run_dos(
+    nodes: &[ContactInfo],
+    iterations: usize,
+    entrypoint_addr: SocketAddr,
+    data_type: String,
+    data_size: usize,
+    mode: String,
+) {
+    let mut target = None;
+    for node in nodes {
+        if node.gossip == entrypoint_addr {
+            target = match mode.as_str() {
+                "gossip" => Some(node.gossip),
+                "tvu" => Some(node.tvu),
+                "tvu_forwards" => Some(node.tvu_forwards),
+                "tpu" => Some(node.tpu),
+                "tpu_forwards" => Some(node.tpu_forwards),
+                "repair" => Some(node.repair),
+                "serve_repair" => Some(node.serve_repair),
+                &_ => panic!("Unknown mode"),
+            };
+            break;
+        }
+    }
+    let target = target.expect("should have target");
+
+    info!("Targetting {}", target);
+    let socket = UdpSocket::bind("0.0.0.0:0").unwrap();
+
+    let mut data = Vec::new();
+
+    let source = thread_rng().gen_range(0, nodes.len());
+    let mut contact = nodes[source].clone();
+    contact.id = Pubkey::new_rand();
+    match data_type.as_str() {
+        "repair_highest" => {
+            let slot = 100;
+            let req = RepairProtocol::WindowIndex(contact, slot, 0);
+            data = bincode::serialize(&req).unwrap();
+        }
+        "repair_shred" => {
+            let slot = 100;
+            let req = RepairProtocol::HighestWindowIndex(contact, slot, 0);
+            data = bincode::serialize(&req).unwrap();
+        }
+        "repair_orphan" => {
+            let slot = 100;
+            let req = RepairProtocol::Orphan(contact, slot);
+            data = bincode::serialize(&req).unwrap();
+        }
+        "random" => {
+            data.resize(data_size, 0);
+        }
+        &_ => {
+            panic!("unknown data type");
+        }
+    }
+
+    let mut last_log = Instant::now();
+    let mut count = 0;
+    let mut error_count = 0;
+    loop {
+        if data_type == "random" {
+            thread_rng().fill(&mut data[..]);
+        }
+        let res = socket.send_to(&data, target);
+        count += 1;
+        if res.is_err() {
+            error_count += 1;
+        }
+        if last_log.elapsed().as_secs() > 5 {
+            info!("count: {} errors: {}", count, error_count);
+            last_log = Instant::now();
+            count = 0;
+        }
+        if iterations != 0 && count >= iterations {
+            break;
+        }
+    }
+}
+
 fn main() {
     solana_logger::setup();
     let matches = App::new("crate")
@@ -18,13 +102,6 @@ fn main() {
                 .value_name("HOST:PORT")
                 .help("Gossip entrypoint address. Usually <ip>:8001"),
         )
-        .arg(
-            Arg::with_name("num_nodes")
-                .long("num_nodes")
-                .takes_value(true)
-                .value_name("NUM")
-                .help("Number of gossip nodes to look for."),
-        )
         .arg(
             Arg::with_name("mode")
                 .long("mode")
@@ -44,10 +121,11 @@ fn main() {
                 .help("Size of packet to dos with."),
         )
         .arg(
-            Arg::with_name("random_data")
+            Arg::with_name("data_type")
-                .long("random_data")
+                .long("data_type")
-                .takes_value(false)
+                .takes_value(true)
-                .help("Use random data for each iteration."),
+                .value_name("DATA_TYPE")
+                .help("Type of data to send."),
         )
         .get_matches();
 
@@ -58,58 +136,64 @@ fn main() {
             exit(1)
         });
     }
-    let num_nodes = value_t!(matches, "num_nodes", usize).unwrap_or(1);
     let data_size = value_t!(matches, "data_size", usize).unwrap_or(128);
-    let random_data = matches.is_present("random_data");
 
     let mode = value_t_or_exit!(matches, "mode", String);
+    let data_type = value_t_or_exit!(matches, "data_type", String);
 
     info!("Finding cluster entry: {:?}", entrypoint_addr);
-    let (nodes, _archivers) = discover_cluster(&entrypoint_addr, num_nodes).unwrap_or_else(|err| {
+    let (nodes, _archivers) = discover(
-        eprintln!("Failed to discover {} nodes: {:?}", num_nodes, err);
+        Some(&entrypoint_addr),
+        None,
+        Some(60),
+        None,
+        Some(&entrypoint_addr),
+        None,
+    )
+    .unwrap_or_else(|err| {
+        eprintln!("Failed to discover {} node: {:?}", entrypoint_addr, err);
         exit(1);
     });
 
     info!("done found {} nodes", nodes.len());
 
-    let mut target = None;
+    run_dos(&nodes, 0, entrypoint_addr, data_type, data_size, mode);
-    for node in &nodes {
+}
-        if node.gossip == entrypoint_addr {
-            target = match mode.as_str() {
-                "gossip" => Some(node.gossip),
-                "tvu" => Some(node.tvu),
-                "tvu_forwards" => Some(node.tvu_forwards),
-                "tpu" => Some(node.tpu),
-                "tpu_forwards" => Some(node.tpu_forwards),
-                "repair" => Some(node.repair),
-                "serve_repair" => Some(node.serve_repair),
-                &_ => panic!("Unknown mode"),
-            };
-            break;
-        }
-    }
-    let target = target.expect("should have target");
 
-    info!("Targetting {}", target);
+#[cfg(test)]
-    let socket = UdpSocket::bind("0.0.0.0:0").unwrap();
+pub mod test {
-    let mut data = vec![0u8; data_size];
+    use super::*;
-    thread_rng().fill(&mut data[..]);
+    use solana_sdk::timing::timestamp;
-    let mut last_log = Instant::now();
+
-    let mut count = 0;
+    #[test]
-    let mut error_count = 0;
+    fn test_dos() {
-    loop {
+        let nodes = [ContactInfo::new_localhost(&Pubkey::new_rand(), timestamp())];
-        if random_data {
+        let entrypoint_addr = nodes[0].gossip.clone();
-            thread_rng().fill(&mut data[..]);
+        run_dos(
-        }
+            &nodes,
-        let res = socket.send_to(&data, target);
+            1,
-        count += 1;
+            entrypoint_addr,
-        if res.is_err() {
+            "random".to_string(),
-            error_count += 1;
+            10,
-        }
+            "tvu".to_string(),
-        if last_log.elapsed().as_secs() > 5 {
+        );
-            info!("count: {} errors: {}", count, error_count);
+
-            last_log = Instant::now();
+        run_dos(
-            count = 0;
+            &nodes,
-        }
+            1,
+            entrypoint_addr,
+            "repair_highest".to_string(),
+            10,
+            "repair".to_string(),
+        );
+
+        run_dos(
+            &nodes,
+            1,
+            entrypoint_addr,
+            "repair_shred".to_string(),
+            10,
+            "serve_repair".to_string(),
+        );
     }
 }