Initial commit

src/main.rs

@@ -0,0 +1,122 @@
+#![feature(let_chains)]
+
+mod client;
+mod record;
+mod server;
+
+use record::Records;
+
+use argp::FromArgs;
+use static_cell::StaticCell;
+use tokio::sync::oneshot;
+
+/// 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,
+}
+
+/// 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 TlsMode {
+	None,
+	Client,
+	Server,
+	Both,
+}
+
+static RECORDS: StaticCell<Records> = StaticCell::new();
+
+#[tokio::main]
+async fn main() {
+	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 records = RECORDS.init(record::read_record_file(&opt.record_file));
+			let (sync_sender, sync_receiver) = oneshot::channel();
+			let client = tokio::spawn(client::play(
+				records,
+				tls_mode,
+				(subopt.forward_addr, subopt.forward_port),
+				sync_receiver,
+				subopt.repeat,
+			));
+			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);
+		}
+	}
+}