netreplay/src/main.rs

mod client;
mod http;
mod record;
mod server;

use record::Records;

use argp::FromArgs;
use static_cell::StaticCell;
use tokio::sync::oneshot;
use tokio_rustls::rustls::crypto::CryptoProvider;

/// 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
	Play(OptPlay),
	/// Print records
	Print(OptPrint),
	/// Record traffic
	Record(OptRecord),
}

/// Replay from records
#[derive(FromArgs)]
#[argp(subcommand, name = "play")]
struct OptPlay {
	/// Connect to address
	#[argp(positional)]
	forward_addr: String,
	/// Connect to port
	#[argp(positional)]
	forward_port: u16,
	/// Listen to port
	#[argp(positional)]
	listen_port: u16,
	/// Path to PEM certificates and keys
	#[argp(positional)]
	certs: String,
	/// Where to use TLS
	#[argp(positional)]
	tls: String,
	/// Repeat N times
	#[argp(option, short = 'r', default = "1")]
	repeat: u32,
	/// Only play this record
	#[argp(option)]
	record: Option<u64>,
	/// Only run these parts
	#[argp(option, default = "String::from(\"both\")")]
	run: String,
}

/// Print records
#[derive(FromArgs)]
#[argp(subcommand, name = "print")]
struct OptPrint {
	/// Print packets
	#[argp(switch, short = 'p')]
	packets: bool,
}

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

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum RunMode {
	Client,
	Server,
	Both,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TlsMode {
	None,
	Client,
	Server,
	Both,
}

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::Play(subopt) => {
			let tls_mode = match subopt.tls.as_str() {
				"none" => TlsMode::None,
				"client" => TlsMode::Client,
				"server" => TlsMode::Server,
				"both" => TlsMode::Both,
				_ => panic!("TLS mode must be one of none,client,server,both."),
			};
			let run_mode = match subopt.run.as_str() {
				"client" => RunMode::Client,
				"server" => RunMode::Server,
				"both" => RunMode::Both,
				_ => panic!("run mode must be one of client,server,both."),
			};
			let records = RECORDS.init(record::read_record_file(&opt.record_file));

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

			let mut ciphers: Option<Vec<String>> = None;
			let mut kexes: Option<Vec<String>> = None;
			for (var, val) in std::env::vars() {
				match var.as_str() {
					"CIPHERS" => ciphers = Some(val.split(',').map(str::to_string).collect()),
					"KEXES" => kexes = Some(val.split(',').map(str::to_string).collect()),
					_ => {}
				}
			}
			let mut prov = tokio_rustls::rustls::crypto::aws_lc_rs::default_provider();
			if let Some(ciphers) = ciphers {
				prov.cipher_suites.clear();
				for cipher in ciphers {
					match cipher.as_str() {
						"AES_256_GCM_SHA384" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS13_AES_256_GCM_SHA384),
						"AES_128_GCM_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS13_AES_128_GCM_SHA256),
						"CHACHA20_POLY1305_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS13_CHACHA20_POLY1305_SHA256),
						"ECDHE_ECDSA_WITH_AES_256_GCM_SHA384" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384),
						"ECDHE_ECDSA_WITH_AES_128_GCM_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256),
						"ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256),
						"ECDHE_RSA_WITH_AES_256_GCM_SHA384" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384),
						"ECDHE_RSA_WITH_AES_128_GCM_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256),
						"ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256" => prov
							.cipher_suites
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::cipher_suite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256),
						other => {
							println!("Unknown cipher `{other}`")
						}
					}
				}
			}
			if let Some(kexes) = kexes {
				prov.kx_groups.clear();
				for kex in kexes {
					match kex.as_str() {
						"X25519" => prov
							.kx_groups
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::kx_group::X25519),
						"SECP256R1" => prov
							.kx_groups
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::kx_group::SECP256R1),
						"SECP384R1" => prov
							.kx_groups
							.push(tokio_rustls::rustls::crypto::aws_lc_rs::kx_group::SECP384R1),
						other => {
							println!("Unknown kex `{other}`")
						}
					}
				}
			}
			CryptoProvider::install_default(prov).unwrap();

			let (sync_sender, sync_receiver) = oneshot::channel();
			//console_subscriber::init();
			let client = tokio::spawn({
				let records = &*records;
				async move {
					if run_mode == RunMode::Both || run_mode == RunMode::Client {
						client::play(
							records,
							tls_mode,
							(subopt.forward_addr, subopt.forward_port),
							sync_receiver,
							subopt.repeat,
						)
						.await;
					} else {
						std::future::pending().await
					}
				}
			});
			if run_mode == RunMode::Both || run_mode == RunMode::Server {
				server::play(
					records,
					tls_mode,
					&subopt.certs,
					("0.0.0.0", subopt.listen_port),
					sync_sender,
				)
				.await;
			}
			client.await.unwrap();
		}
		Subcommand::Print(subopt) => {
			let records = record::read_record_file(&opt.record_file);
			record::print_records(&records, subopt.packets);
		}
		Subcommand::Record(_subopt) => {
			record::make_record(&opt.record_file);
		}
	}
}