From 9b095a7938599f13788ff0cd726330eaa5877ed3 Mon Sep 17 00:00:00 2001
From: Kris Kwiatkowski <kris@cloudflare.com>
Date: Fri, 26 Jul 2019 14:02:54 +0100
Subject: [PATCH] process in multiple threads

---
 rawhammer/Cargo.toml  |  2 --
 rawhammer/src/main.rs | 81 ++++++++++++++++++++++---------------------
 2 files changed, 42 insertions(+), 41 deletions(-)

diff --git a/rawhammer/Cargo.toml b/rawhammer/Cargo.toml
index fbcaa69..7ce5a12 100644
--- a/rawhammer/Cargo.toml
+++ b/rawhammer/Cargo.toml
@@ -7,6 +7,4 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-ctrlc = "3.1.3"
-thread-control = "0.1.2"
 clap = "2.33.0"
diff --git a/rawhammer/src/main.rs b/rawhammer/src/main.rs
index dffb895..674b4c3 100644
--- a/rawhammer/src/main.rs
+++ b/rawhammer/src/main.rs
@@ -3,8 +3,6 @@ use std::net::*;
 use std::time::Duration;
 use std::thread;
 use std::time;
-use ctrlc;
-use thread_control::*;
 use clap::{Arg, App};
 
 static CH_X25519: &'static [u8] = b"\
@@ -204,15 +202,6 @@ fn sleep_no_comment(millis: u64) {
 }
 
 fn main() -> std::io::Result<()> {
-    // Condition variables to control threads
-    let (flag_t1, control_t1) = make_pair();
-
-    ctrlc::set_handler(move || {
-        println!("Stopping...");
-        control_t1.stop();
-    }).expect("Error setting Ctrl-C handler");
-
-
     let matches = App::new("rawhammer")
         .version("0.0.1")
         .author("Kris Kwiatkowski <kris@amongbytes.com>")
@@ -241,6 +230,12 @@ fn main() -> std::io::Result<()> {
                  .takes_value(true)
                  .required(false)
                  .help("Number of connections sent per second (default 100)"))
+        .arg(Arg::with_name("threads")
+                 .short("s")
+                 .long("sending_threads")
+                 .takes_value(true)
+                 .required(false)
+                 .help("Number of sending threads (default 1)"))
         .get_matches();
 
     let sock_addr: SocketAddr = matches
@@ -252,40 +247,48 @@ fn main() -> std::io::Result<()> {
         .parse().unwrap(), 0);
     let msg_per_sec: usize =  matches.value_of("frequency").unwrap_or("100").parse().unwrap();
     let ch_case: usize = matches.value_of("case").unwrap().parse().unwrap();
+    let sending_threads: usize =  matches.value_of("threads").unwrap_or("1").parse().unwrap();
 
-    println!("Sending to {0}, with freq {1} per sec.", sock_addr, msg_per_sec);
+    println!("Sending to {0}, with freq {1} per sec on {2} threads.", sock_addr, msg_per_sec, sending_threads);
 
     // Start a thread which sends CH
-    let _t1 = thread::spawn(move || {
-        // calculates number of sent client hellos
-        let mut ch_sent = 0;
+    let mut threads = Vec::new();
+    for _ in 0..sending_threads {
+        // Condition variables to control threads
+        let t = thread::spawn(move || {
+            // calculates number of sent client hellos
+            let mut ch_sent = 0;
 
-        // Sent until CTRL+C
-        while flag_t1.is_alive() {
-            // try connect
-            let sock = TcpStream::connect_timeout(&sock_addr, dur);
-            if sock.is_err() {
-                sleep(10, "Can't connect, waiting...");
-                continue;
-            }
-            let mut s = sock.unwrap();
-            match ch_case {
-                1 => s.write(&CH_NTRU).unwrap(),
-                2 => s.write(&CH_SIKE).unwrap(),
-                3 => s.write(&CH_X25519).unwrap(),
-                4 => s.write(&CH_P521).unwrap(),
-                _ => panic!("Unknown case"),
-            };
-            s.shutdown(Shutdown::Both).expect("shutdown call failed");
-            sleep_no_comment(1000/msg_per_sec as u64);
-            ch_sent += 1;
-            if ch_sent % 100 == 0 {
-                println!("nb of queries sent: {}", ch_sent);
+            // Sent until CTRL+C
+            loop {
+                // try connect
+                let sock = TcpStream::connect_timeout(&sock_addr, dur);
+                if sock.is_err() {
+                    sleep(10, "Can't connect, waiting...");
+                    continue;
+                }
+                let mut s = sock.unwrap();
+                match ch_case {
+                    1 => s.write(&CH_NTRU).unwrap(),
+                    2 => s.write(&CH_SIKE).unwrap(),
+                    3 => s.write(&CH_X25519).unwrap(),
+                    4 => s.write(&CH_P521).unwrap(),
+                    _ => panic!("Unknown case"),
+                };
+                s.shutdown(Shutdown::Both).expect("shutdown call failed");
+                sleep_no_comment(1000/msg_per_sec as u64);
+                ch_sent += 1;
+                if ch_sent % 100 == 0 {
+                    println!("nb of queries sent: {}", ch_sent);
+                }
             }
-        }
-    });
+        });
+        threads.push(t);
+    }
 
     // wait for TX thread to finish
-    _t1.join().unwrap();
+    for t in threads {
+        t.join().unwrap();
+    }
     Ok(())
 }