netreplay/src/main.rs

#![feature(ascii_char)]

mod client;
mod codec;
mod record;
mod server;
mod util;

use record::Records;

use argp::FromArgs;
use static_cell::StaticCell;

/// Play recorded requests and responses
#[derive(FromArgs)]
struct Opt {
	/// Path to record file
	#[argp(positional)]
	record_file: String,
	#[argp(subcommand)]
	subcommand: Subcommand,
}

#[derive(FromArgs)]
#[argp(subcommand)]
enum Subcommand {
	/// Replay from records (client)
	Client(OptClient),
	/// Replay from records (server)
	Server(OptServer),
	/// Print records
	Print(OptPrint),
	/// Record traffic
	#[cfg(feature = "record")]
	Record(OptRecord),
	/// Remove record
	Remove(OptRemove),
	/// Write test record
	Test(OptTest),
}

/// Replay from records
#[derive(FromArgs)]
#[argp(subcommand, name = "client")]
struct OptClient {
	/// Connect to address
	#[argp(positional)]
	connect_addr: String,
	/// Connect to port
	#[argp(positional)]
	connect_port: u16,
	/// Whether to use TLS
	#[argp(switch, long = "tls")]
	tls: bool,
	/// Repeat N times
	#[argp(option, short = 'r', default = "1")]
	repeat: u32,
	/// UDP end notification will be sent to this address:port
	#[argp(option, short = 'n')]
	notify_addr: Option<String>,
	/// Only play this record
	#[argp(option)]
	record: Option<u64>,
	/// Path to PEM certificates (if not provided, use system's certificates)
	#[argp(option, short = 'c')]
	certs: Option<String>,
	/// Do not verify certificates
	#[argp(switch, short = 's')]
	skip_verif: bool,
	/// Maximum number of concurrent clients
	#[argp(option, default = "16")]
	concurrency: usize,
	/// Print debug info
	#[argp(switch, short = 'd')]
	debug: bool,
}

/// Replay from records
#[derive(FromArgs)]
#[argp(subcommand, name = "server")]
struct OptServer {
	/// Listen to port
	#[argp(positional)]
	listen_port: u16,
	/// Path to PEM certificates and keys
	#[argp(positional)]
	certs: String,
	/// Whether to use TLS
	#[argp(switch, long = "tls")]
	tls: bool,
	/// Print debug info
	#[argp(switch, short = 'd')]
	debug: bool,
}

/// Print records
#[derive(FromArgs)]
#[argp(subcommand, name = "print")]
struct OptPrint {
	/// Record number
	#[argp(option, short = 'n')]
	number: Option<u64>,
}

/// Record traffic
#[cfg(feature = "record")]
#[derive(FromArgs)]
#[argp(subcommand, name = "record")]
struct OptRecord {}

/// Record traffic
#[derive(FromArgs)]
#[argp(subcommand, name = "test")]
struct OptTest {}

/// Copy record but removing one connection id
#[derive(FromArgs)]
#[argp(subcommand, name = "remove")]
struct OptRemove {
	/// Output path
	#[argp(positional)]
	output: String,
	/// Record number to remove
	#[argp(positional)]
	record_number: u64,
	/// Packet number to remove
	#[argp(positional)]
	packet_number: usize,
}

static RECORDS: StaticCell<Records> = StaticCell::new();

#[tokio::main]
async fn main() {
	env_logger::init();
	let opt: Opt = argp::parse_args_or_exit(argp::DEFAULT);

	match opt.subcommand {
		Subcommand::Client(subopt) => {
			let records = RECORDS.init(record::read_record_file(&opt.record_file));

			if let Some(only_record) = subopt.record {
				records.retain(|id, _| *id == only_record);
			}

			util::init_provider();

			//console_subscriber::init();
			client::play(
				records,
				subopt.tls,
				(subopt.connect_addr, subopt.connect_port),
				subopt.repeat,
				subopt.certs.as_deref(),
				subopt.skip_verif,
				subopt.concurrency,
				subopt.notify_addr.as_deref(),
				subopt.debug,
			)
			.await;
		}
		Subcommand::Server(subopt) => {
			let records = RECORDS.init(record::read_record_file(&opt.record_file));

			util::init_provider();

			//console_subscriber::init();
			server::play(
				records,
				subopt.tls,
				&subopt.certs,
				("0.0.0.0", subopt.listen_port),
				subopt.debug,
			)
			.await;
		}
		Subcommand::Print(subopt) => {
			let records = record::read_record_file(&opt.record_file);
			record::print_records(&records, subopt.number);
		}
		#[cfg(feature = "record")]
		Subcommand::Record(_subopt) => {
			record::make_record(&opt.record_file);
		}
		Subcommand::Remove(subopt) => {
			record::remove_record(
				&opt.record_file,
				&subopt.output,
				subopt.record_number,
				subopt.packet_number,
			);
		}
		Subcommand::Test(_subopt) => {
			record::make_test_record(&opt.record_file);
		}
	}
}