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,8 +18,16 @@ fn run_dos(
     data_type: String,
     data_size: usize,
     mode: String,
+    data_input: Option<String>,
 ) {
     let mut target = None;
+    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() {
@@ -28,11 +38,16 @@ fn run_dos(
                     "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");
 
     info!("Targetting {}", target);
@@ -40,6 +55,7 @@ fn run_dos(
 
     let mut data = Vec::new();
 
+    if !nodes.is_empty() {
         let source = thread_rng().gen_range(0, nodes.len());
         let mut contact = nodes[source].clone();
         contact.id = Pubkey::new_rand();
@@ -62,23 +78,51 @@ fn run_dos(
             "random" => {
                 data.resize(data_size, 0);
             }
+            "get_account_info" => {}
+            "get_program_accounts" => {}
             &_ => {
                 panic!("unknown data type");
             }
         }
+    }
 
     let mut last_log = Instant::now();
     let mut count = 0;
     let mut error_count = 0;
     loop {
+        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);
-        count += 1;
             if res.is_err() {
                 error_count += 1;
             }
+        }
+        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,12 +207,16 @@ 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();
 
+    let mut nodes = vec![];
+    if !skip_gossip {
         info!("Finding cluster entry: {:?}", entrypoint_addr);
-    let (nodes, _validators) = discover(
+        let (gossip_nodes, _validators) = discover(
             None,
             Some(&entrypoint_addr),
             None,
@@ -162,10 +230,20 @@ fn main() {
             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,
         );
     }
 }