Add feature to RPC dos (#12119)

Cargo.lock

@@ -3678,6 +3678,7 @@ dependencies = [
  "rand 0.7.3",
  "rayon",
  "solana-clap-utils",
+ "solana-client",
  "solana-core",
  "solana-ledger",
  "solana-logger 1.4.0",

dos/Cargo.toml

@@ -21,6 +21,7 @@ solana-net-utils = { path = "../net-utils", version = "1.4.0" }
 solana-runtime = { path = "../runtime", version = "1.4.0" }
 solana-sdk = { path = "../sdk", version = "1.4.0" }
 solana-version = { path = "../version", version = "1.4.0" }
+solana-client = { path = "../client", version = "1.4.0" }
 
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]

dos/src/main.rs

@@ -1,12 +1,14 @@
 use clap::{crate_description, crate_name, value_t, value_t_or_exit, App, Arg};
 use log::*;
 use rand::{thread_rng, Rng};
+use solana_client::rpc_client::RpcClient;
 use solana_core::{
     contact_info::ContactInfo, gossip_service::discover, serve_repair::RepairProtocol,
 };
 use solana_sdk::pubkey::Pubkey;
 use std::net::{SocketAddr, UdpSocket};
 use std::process::exit;
+use std::str::FromStr;
 use std::time::Instant;
 
 fn run_dos(
@@ -16,21 +18,34 @@ fn run_dos(
     data_type: String,
     data_size: usize,
     mode: String,
+    data_input: Option<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 mut rpc_client = None;
+    if nodes.is_empty() {
+        if mode == "rpc" {
+            rpc_client = Some(RpcClient::new_socket(entrypoint_addr));
+        }
+        target = Some(entrypoint_addr);
+    } else {
+        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),
+                    "rpc" => {
+                        rpc_client = Some(RpcClient::new_socket(node.rpc));
+                        None
+                    }
+                    &_ => panic!("Unknown mode"),
+                };
+                break;
+            }
         }
     }
     let target = target.expect("should have target");
@@ -40,30 +55,34 @@ fn run_dos(
 
     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::WindowIndexWithNonce(contact, slot, 0, 0);
-            data = bincode::serialize(&req).unwrap();
-        }
-        "repair_shred" => {
-            let slot = 100;
-            let req = RepairProtocol::HighestWindowIndexWithNonce(contact, slot, 0, 0);
-            data = bincode::serialize(&req).unwrap();
-        }
-        "repair_orphan" => {
-            let slot = 100;
-            let req = RepairProtocol::OrphanWithNonce(contact, slot, 0);
-            data = bincode::serialize(&req).unwrap();
-        }
-        "random" => {
-            data.resize(data_size, 0);
-        }
-        &_ => {
-            panic!("unknown data type");
+    if !nodes.is_empty() {
+        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::WindowIndexWithNonce(contact, slot, 0, 0);
+                data = bincode::serialize(&req).unwrap();
+            }
+            "repair_shred" => {
+                let slot = 100;
+                let req = RepairProtocol::HighestWindowIndexWithNonce(contact, slot, 0, 0);
+                data = bincode::serialize(&req).unwrap();
+            }
+            "repair_orphan" => {
+                let slot = 100;
+                let req = RepairProtocol::OrphanWithNonce(contact, slot, 0);
+                data = bincode::serialize(&req).unwrap();
+            }
+            "random" => {
+                data.resize(data_size, 0);
+            }
+            "get_account_info" => {}
+            "get_program_accounts" => {}
+            &_ => {
+                panic!("unknown data type");
+            }
         }
     }
 
@@ -71,14 +90,39 @@ fn run_dos(
     let mut count = 0;
     let mut error_count = 0;
     loop {
-        if data_type == "random" {
-            thread_rng().fill(&mut data[..]);
+        if mode == "rpc" {
+            match data_type.as_str() {
+                "get_account_info" => {
+                    let res = rpc_client
+                        .as_ref()
+                        .unwrap()
+                        .get_account(&Pubkey::from_str(&data_input.as_ref().unwrap()).unwrap());
+                    if res.is_err() {
+                        error_count += 1;
+                    }
+                }
+                "get_program_accounts" => {
+                    let res = rpc_client.as_ref().unwrap().get_program_accounts(
+                        &Pubkey::from_str(&data_input.as_ref().unwrap()).unwrap(),
+                    );
+                    if res.is_err() {
+                        error_count += 1;
+                    }
+                }
+                &_ => {
+                    panic!("unsupported data type");
+                }
+            }
+        } else {
+            if data_type == "random" {
+                thread_rng().fill(&mut data[..]);
+            }
+            let res = socket.send_to(&data, target);
+            if res.is_err() {
+                error_count += 1;
+            }
         }
-        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();
@@ -115,6 +159,7 @@ fn main() {
                     "tpu_forwards",
                     "repair",
                     "serve_repair",
+                    "rpc",
                 ])
                 .help("Interface to DoS"),
         )
@@ -130,9 +175,28 @@ fn main() {
                 .long("data-type")
                 .takes_value(true)
                 .value_name("TYPE")
-                .possible_values(&["repair_highest", "repair_shred", "repair_orphan", "random"])
+                .possible_values(&[
+                    "repair_highest",
+                    "repair_shred",
+                    "repair_orphan",
+                    "random",
+                    "get_account_info",
+                    "get_program_accounts",
+                ])
                 .help("Type of data to send"),
         )
+        .arg(
+            Arg::with_name("data_input")
+                .long("data-input")
+                .takes_value(true)
+                .value_name("TYPE")
+                .help("Data to send"),
+        )
+        .arg(
+            Arg::with_name("skip_gossip")
+                .long("skip-gossip")
+                .help("Just use entrypoint address directly"),
+        )
         .get_matches();
 
     let mut entrypoint_addr = SocketAddr::from(([127, 0, 0, 1], 8001));
@@ -143,29 +207,43 @@ fn main() {
         });
     }
     let data_size = value_t!(matches, "data_size", usize).unwrap_or(128);
+    let skip_gossip = matches.is_present("skip_gossip");
 
     let mode = value_t_or_exit!(matches, "mode", String);
     let data_type = value_t_or_exit!(matches, "data_type", String);
+    let data_input = value_t!(matches, "data_input", String).ok();
 
-    info!("Finding cluster entry: {:?}", entrypoint_addr);
-    let (nodes, _validators) = discover(
-        None,
-        Some(&entrypoint_addr),
-        None,
-        Some(60),
-        None,
-        Some(&entrypoint_addr),
-        None,
-        0,
-    )
-    .unwrap_or_else(|err| {
-        eprintln!("Failed to discover {} node: {:?}", entrypoint_addr, err);
-        exit(1);
-    });
+    let mut nodes = vec![];
+    if !skip_gossip {
+        info!("Finding cluster entry: {:?}", entrypoint_addr);
+        let (gossip_nodes, _validators) = discover(
+            None,
+            Some(&entrypoint_addr),
+            None,
+            Some(60),
+            None,
+            Some(&entrypoint_addr),
+            None,
+            0,
+        )
+        .unwrap_or_else(|err| {
+            eprintln!("Failed to discover {} node: {:?}", entrypoint_addr, err);
+            exit(1);
+        });
+        nodes = gossip_nodes;
+    }
 
     info!("done found {} nodes", nodes.len());
 
-    run_dos(&nodes, 0, entrypoint_addr, data_type, data_size, mode);
+    run_dos(
+        &nodes,
+        0,
+        entrypoint_addr,
+        data_type,
+        data_size,
+        mode,
+        data_input,
+    );
 }
 
 #[cfg(test)]
@@ -184,6 +262,7 @@ pub mod test {
             "random".to_string(),
             10,
             "tvu".to_string(),
+            None,
         );
 
         run_dos(
@@ -193,6 +272,7 @@ pub mod test {
             "repair_highest".to_string(),
             10,
             "repair".to_string(),
+            None,
         );
 
         run_dos(
@@ -202,6 +282,7 @@ pub mod test {
             "repair_shred".to_string(),
             10,
             "serve_repair".to_string(),
+            None,
         );
     }
 }