diff --git a/Cargo.lock b/Cargo.lock index d105c5f7..4808748a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6005,6 +6005,8 @@ dependencies = [ "bytes", "clap", "futures", + "rand 0.10.2", + "tailscale", "tokio", "tracing", "ts_cli_util", diff --git a/ts_devtools/Cargo.toml b/ts_devtools/Cargo.toml index a49ee220..e0f311a8 100644 --- a/ts_devtools/Cargo.toml +++ b/ts_devtools/Cargo.toml @@ -17,7 +17,9 @@ clap.workspace = true tokio = { workspace = true, features = ["full"] } futures.workspace = true tracing.workspace = true +rand.workspace = true +tailscale.workspace = true ts_cli_util.workspace = true ts_keys.workspace = true ts_packetfilter.workspace = true diff --git a/ts_devtools/src/bin/throughput.rs b/ts_devtools/src/bin/throughput.rs new file mode 100644 index 00000000..ada0d748 --- /dev/null +++ b/ts_devtools/src/bin/throughput.rs @@ -0,0 +1,81 @@ +//! TCP server which sinks all incoming traffic and writes out random data as fast as +//! possible. + +use core::{ + convert::Infallible, + error::Error, + pin::Pin, + task::{Context, Poll}, +}; + +use clap::Parser; +use rand::Rng; +use tokio::io::ReadBuf; + +#[derive(Debug, Clone, clap::Parser)] +struct Args { + #[clap(flatten)] + common: ts_cli_util::CommonArgs, +} + +#[tokio::main] +async fn main() -> Result> { + ts_cli_util::init_tracing(); + let args = Args::parse(); + + let config = args.common.config().await?; + + let dev = tailscale::Device::new(&config, None).await?; + let listener = dev + .tcp_listen((dev.ipv4_addr().await?, 1234).into()) + .await?; + + tracing::info!(endpoint = %listener.local_addr(), "tcp socket listening"); + + loop { + let mut sock = match listener.accept().await { + Ok(sock) => sock, + Err(e) => { + tracing::error!(error = %e, "accepting connection"); + continue; + } + }; + + tracing::info!(remote = %sock.remote_addr(), "accept"); + + tokio::task::spawn(async move { + let mut dst = tokio::io::join(AsyncRng, tokio::io::sink()); + + if let Err(e) = tokio::io::copy_bidirectional(&mut sock, &mut dst).await { + tracing::error!(remote = %sock.remote_addr(), error = %e); + } else { + tracing::info!(remote = %sock.remote_addr(), "close"); + } + }); + } +} + +/// [`AsyncRead`][tokio::io::AsyncRead] which always immediately and completely fills any +/// buffer with random data using [`rand::rng`]. +/// +/// This technically isn't holding tokio correctly because it can potentially do a lot of +/// CPU-bound work inside the poll, but for this informal benchmark it hasn't proven to be +/// a bottleneck. +#[derive(Copy, Clone)] +struct AsyncRng; + +impl tokio::io::AsyncRead for AsyncRng { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let unfilled = buf.initialize_unfilled(); + rand::rng().fill_bytes(unfilled); + let len = unfilled.len(); + + buf.advance(len); + + Poll::Ready(Ok(())) + } +} diff --git a/ts_runtime/src/lib.rs b/ts_runtime/src/lib.rs index 9fd6b0ba..d483a662 100644 --- a/ts_runtime/src/lib.rs +++ b/ts_runtime/src/lib.rs @@ -115,7 +115,11 @@ impl kameo::Actor for Runtime { &slf, ( env.clone(), - Default::default(), + netstack::netcore::Config { + tcp_buffer_size: 64 * 1024, + command_channel_capacity: Some(128), + ..Default::default() + }, netstack_id, netstack_up, Arc::new(Mutex::new(netstack_down)),