Something went wrong. Try again.
noq, a QUIC implementation in Rust
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295use core::str;use std::{ convert::TryInto, net::{IpAddr, Ipv6Addr, SocketAddr}, num::ParseIntError, str::FromStr, sync::Arc,};
use anyhow::{Context, Result};use bytes::Bytes;use clap::Parser;use noq::crypto::rustls::QuicClientConfig;use rustls::{ RootCertStore, pki_types::{CertificateDer, PrivateKeyDer},};use tokio::runtime::{Builder, Runtime};use tracing::trace;
pub mod stats;
pub fn configure_tracing_subscriber() { tracing::subscriber::set_global_default( tracing_subscriber::FmtSubscriber::builder() .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) .finish(), ) .unwrap();}
/// Creates a server endpoint which runs on the given runtimepub fn server_endpoint( rt: &Runtime, cert: CertificateDer<'static>, key: PrivateKeyDer<'static>, opt: &Opt,) -> (SocketAddr, noq::Endpoint) { let cert_chain = vec![cert]; let mut server_config = noq::ServerConfig::with_single_cert(cert_chain, key).unwrap(); server_config.transport = Arc::new(transport_config(opt));
let endpoint = { let _guard = rt.enter(); noq::Endpoint::server( server_config, SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 0), ) .unwrap() }; let server_addr = endpoint.local_addr().unwrap(); (server_addr, endpoint)}
/// Create a client endpoint and client connectionpub async fn connect_client( server_addr: SocketAddr, server_cert: CertificateDer<'_>, opt: Opt,) -> Result<(noq::Endpoint, noq::Connection)> { let endpoint = noq::Endpoint::client(SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 0)).unwrap();
let mut roots = RootCertStore::empty(); roots.add(server_cert)?;
let default_provider = rustls::crypto::ring::default_provider(); let provider = rustls::crypto::CryptoProvider { cipher_suites: vec![opt.cipher.as_rustls()], ..default_provider };
let crypto = rustls::ClientConfig::builder_with_provider(provider.into()) .with_protocol_versions(&[&rustls::version::TLS13]) .unwrap() .with_root_certificates(roots) .with_no_client_auth();
let mut client_config = noq::ClientConfig::new(Arc::new(QuicClientConfig::try_from(crypto)?)); client_config.transport_config(Arc::new(transport_config(&opt)));
let connection = endpoint .connect_with(client_config, server_addr, "localhost") .unwrap() .await .context("unable to connect")?; trace!("connected");
Ok((endpoint, connection))}
pub async fn drain_stream(mut stream: noq::RecvStream, read_unordered: bool) -> Result<usize> { let mut read = 0;
if read_unordered { let mut stream = stream.into_unordered(); while let Some(chunk) = stream.read_chunk(usize::MAX).await? { read += chunk.bytes.len(); } } else { // These are 32 buffers, for reading approximately 32kB at once #[rustfmt::skip] let mut bufs = [ Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(), ];
while let Some(n) = stream.read_many_chunks(&mut bufs[..]).await? { read += bufs.iter().take(n).map(|buf| buf.len()).sum::<usize>(); } }
Ok(read)}
pub async fn send_data_on_stream(stream: &mut noq::SendStream, stream_size: u64) -> Result<()> { const DATA: &[u8] = &[0xAB; 1024 * 1024]; let bytes_data = Bytes::from_static(DATA);
let full_chunks = stream_size / (DATA.len() as u64); let remaining = (stream_size % (DATA.len() as u64)) as usize;
for _ in 0..full_chunks { stream .write_chunk(bytes_data.clone()) .await .context("failed sending data")?; }
if remaining != 0 { stream .write_chunk(bytes_data.slice(0..remaining)) .await .context("failed sending data")?; }
stream.finish().unwrap(); // Wait for stream to close _ = stream.stopped().await;
Ok(())}
pub fn rt(runtime_type: RuntimeType) -> Runtime { match runtime_type { RuntimeType::Tokio => { let counter = std::sync::atomic::AtomicUsize::new(0); Builder::new_multi_thread() .thread_name_fn(move || { format!( "tokio-runtime-{}", counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) ) }) .enable_all() .build() .unwrap() } RuntimeType::TokioCurrentThread => { Builder::new_current_thread().enable_all().build().unwrap() } }}
pub fn transport_config(opt: &Opt) -> noq::TransportConfig { // High stream windows are chosen because the amount of concurrent streams // is configurable as a parameter. let mut config = noq::TransportConfig::default(); config.max_concurrent_uni_streams(opt.max_streams.try_into().unwrap()); config.initial_mtu(opt.initial_mtu);
let mut acks = noq::AckFrequencyConfig::default(); acks.ack_eliciting_threshold(10u32.into()); config.ack_frequency_config(Some(acks));
config}
#[derive(Parser, Debug, Clone, Copy)]#[clap(name = "bulk")]pub struct Opt { /// The total number of clients which should be created #[clap(long = "clients", short = 'c', default_value = "1")] pub clients: usize, /// The total number of streams which should be created #[clap(long = "streams", short = 'n', default_value = "1")] pub streams: usize, /// The amount of concurrent streams which should be used #[clap(long = "max_streams", short = 'm', default_value = "1")] pub max_streams: usize, /// Number of bytes to transmit from server to client /// /// This can use SI suffixes for sizes. For example, 1M will transfer /// 1MiB, 10G will transfer 10GiB. #[clap(long, default_value = "1G", value_parser = parse_byte_size)] pub download_size: u64, /// Number of bytes to transmit from client to server /// /// This can use SI suffixes for sizes. For example, 1M will transfer /// 1MiB, 10G will transfer 10GiB. #[clap(long, default_value = "0", value_parser = parse_byte_size)] pub upload_size: u64, /// Show connection stats the at the end of the benchmark #[clap(long = "stats")] pub stats: bool, /// Whether to use the unordered read API #[clap(long = "unordered")] pub read_unordered: bool, /// Allows to configure the desired cipher suite /// /// Valid options are: aes128, aes256, chacha20 #[clap(long = "cipher", default_value = "aes128")] pub cipher: CipherSuite, /// Starting guess for maximum UDP payload size #[clap(long, default_value = "1200")] pub initial_mtu: u16, /// The runtime type to use #[clap(long, default_value = "tokio")] pub runtime_type: RuntimeType,}
#[derive(Debug, Clone, Copy)]pub enum RuntimeType { Tokio, TokioCurrentThread,}
impl FromStr for RuntimeType { type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> { match s.to_lowercase().as_str() { "tokio" => Ok(Self::Tokio), "tokio-current-thread" => Ok(Self::TokioCurrentThread), _ => Err(anyhow::anyhow!("Unknown runtime type {}", s)), } }}
fn parse_byte_size(s: &str) -> Result<u64, ParseIntError> { let s = s.trim();
let multiplier = match s.chars().last() { Some('T') => 1024 * 1024 * 1024 * 1024, Some('G') => 1024 * 1024 * 1024, Some('M') => 1024 * 1024, Some('k') => 1024, _ => 1, };
let s = match multiplier { 1 => s, _ => &s[..s.len() - 1], };
Ok(u64::from_str(s)? * multiplier)}
#[derive(Debug, PartialEq, Eq, Clone, Copy)]pub enum CipherSuite { Aes128, Aes256, Chacha20,}
impl CipherSuite { pub fn as_rustls(self) -> rustls::SupportedCipherSuite { use rustls::crypto::ring::cipher_suite; match self { Self::Aes128 => cipher_suite::TLS13_AES_128_GCM_SHA256, Self::Aes256 => cipher_suite::TLS13_AES_256_GCM_SHA384, Self::Chacha20 => cipher_suite::TLS13_CHACHA20_POLY1305_SHA256, } }}
impl FromStr for CipherSuite { type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> { match s.to_lowercase().as_str() { "aes128" => Ok(Self::Aes128), "aes256" => Ok(Self::Aes256), "chacha20" => Ok(Self::Chacha20), _ => Err(anyhow::anyhow!("Unknown cipher suite {}", s)), } }}