132 lines
2.7 KiB
Rust
132 lines
2.7 KiB
Rust
#![feature(let_chains)]
|
|
|
|
mod client;
|
|
mod http;
|
|
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,
|
|
/// Only play this record
|
|
#[argp(option)]
|
|
record: Option<u64>,
|
|
}
|
|
|
|
/// 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));
|
|
|
|
if let Some(only_record) = subopt.record {
|
|
records.retain(|id, _| *id == only_record);
|
|
}
|
|
|
|
let (sync_sender, sync_receiver) = oneshot::channel();
|
|
console_subscriber::init();
|
|
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);
|
|
}
|
|
}
|
|
}
|