perf/
server.rs

1use std::{net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
2
3use anyhow::{Context, Result};
4use bytes::Bytes;
5use clap::Parser;
6use quinn::{TokioRuntime, crypto::rustls::QuicServerConfig};
7use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, pem::PemObject};
8use tracing::{debug, error, info};
9
10use crate::{CommonOpt, PERF_CIPHER_SUITES, noprotection::NoProtectionServerConfig};
11
12#[derive(Parser)]
13#[clap(name = "server")]
14pub struct Opt {
15    /// Address to listen on
16    #[clap(long = "listen", default_value = "[::]:4433")]
17    listen: SocketAddr,
18    /// TLS private key in DER format
19    #[clap(short = 'k', long = "key", requires = "cert")]
20    key: Option<PathBuf>,
21    /// TLS certificate in PEM format
22    #[clap(short = 'c', long = "cert", requires = "key")]
23    cert: Option<PathBuf>,
24    /// Common options
25    #[command(flatten)]
26    common: CommonOpt,
27}
28
29pub async fn run(opt: Opt) -> Result<()> {
30    let (key, cert) = match (&opt.key, &opt.cert) {
31        (Some(key), Some(cert)) => (
32            PrivateKeyDer::from_pem_file(key).context("reading private key")?,
33            CertificateDer::pem_file_iter(cert)
34                .context("reading certificate chain file")?
35                .collect::<Result<_, _>>()
36                .context("reading certificate chain")?,
37        ),
38        _ => {
39            let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
40            (
41                PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der()).into(),
42                vec![CertificateDer::from(cert.cert)],
43            )
44        }
45    };
46
47    let default_provider = rustls::crypto::ring::default_provider();
48    let provider = rustls::crypto::CryptoProvider {
49        cipher_suites: PERF_CIPHER_SUITES.into(),
50        ..default_provider
51    };
52
53    let mut crypto = rustls::ServerConfig::builder_with_provider(provider.into())
54        .with_protocol_versions(&[&rustls::version::TLS13])
55        .unwrap()
56        .with_no_client_auth()
57        .with_single_cert(cert, key)
58        .unwrap();
59    crypto.alpn_protocols = vec![b"perf".to_vec()];
60
61    if opt.common.keylog {
62        crypto.key_log = Arc::new(rustls::KeyLogFile::new());
63    }
64
65    let transport = opt.common.build_transport_config(
66        #[cfg(feature = "qlog")]
67        "perf-server",
68    )?;
69
70    let crypto = Arc::new(QuicServerConfig::try_from(crypto)?);
71    let mut config = quinn::ServerConfig::with_crypto(match opt.common.no_protection {
72        true => Arc::new(NoProtectionServerConfig::new(crypto)),
73        false => crypto,
74    });
75    config.transport_config(Arc::new(transport));
76
77    let socket = opt.common.bind_socket(opt.listen)?;
78
79    let mut endpoint_cfg = quinn::EndpointConfig::default();
80    endpoint_cfg.max_udp_payload_size(opt.common.max_udp_payload_size)?;
81
82    let endpoint = quinn::Endpoint::new(endpoint_cfg, Some(config), socket, Arc::new(TokioRuntime))
83        .context("creating endpoint")?;
84
85    info!("listening on {}", endpoint.local_addr().unwrap());
86
87    let opt = Arc::new(opt);
88
89    while let Some(handshake) = endpoint.accept().await {
90        let opt = opt.clone();
91        tokio::spawn(async move {
92            if let Err(e) = handle(handshake, opt).await {
93                error!("connection lost: {:#}", e);
94            }
95        });
96    }
97
98    Ok(())
99}
100
101async fn handle(handshake: quinn::Incoming, opt: Arc<Opt>) -> Result<()> {
102    let connection = handshake.await.context("handshake failed")?;
103
104    debug!("{} connected", connection.remote_address());
105    tokio::try_join!(
106        drive_uni(connection.clone()),
107        drive_bi(connection.clone()),
108        conn_stats(connection, opt)
109    )?;
110    Ok(())
111}
112
113async fn conn_stats(connection: quinn::Connection, opt: Arc<Opt>) -> Result<()> {
114    if opt.common.conn_stats {
115        loop {
116            tokio::time::sleep(Duration::from_secs(2)).await;
117            println!("{:?}\n", connection.stats());
118        }
119    }
120
121    Ok(())
122}
123
124async fn drive_uni(connection: quinn::Connection) -> Result<()> {
125    while let Ok(stream) = connection.accept_uni().await {
126        let connection = connection.clone();
127        tokio::spawn(async move {
128            if let Err(e) = handle_uni(connection, stream).await {
129                error!("request failed: {:#}", e);
130            }
131        });
132    }
133    Ok(())
134}
135
136async fn handle_uni(connection: quinn::Connection, stream: quinn::RecvStream) -> Result<()> {
137    let bytes = read_req(stream).await?;
138    let response = connection.open_uni().await?;
139    respond(bytes, response).await?;
140    Ok(())
141}
142
143async fn drive_bi(connection: quinn::Connection) -> Result<()> {
144    while let Ok((send, recv)) = connection.accept_bi().await {
145        tokio::spawn(async move {
146            if let Err(e) = handle_bi(send, recv).await {
147                error!("request failed: {:#}", e);
148            }
149        });
150    }
151    Ok(())
152}
153
154async fn handle_bi(send: quinn::SendStream, recv: quinn::RecvStream) -> Result<()> {
155    let bytes = read_req(recv).await?;
156    respond(bytes, send).await?;
157    Ok(())
158}
159
160async fn read_req(mut stream: quinn::RecvStream) -> Result<u64> {
161    let mut buf = [0; 8];
162    stream
163        .read_exact(&mut buf)
164        .await
165        .context("reading request")?;
166    let n = u64::from_be_bytes(buf);
167    debug!("got req for {} bytes on {}", n, stream.id());
168    drain_stream(stream).await?;
169    Ok(n)
170}
171
172async fn drain_stream(mut stream: quinn::RecvStream) -> Result<()> {
173    #[rustfmt::skip]
174    let mut bufs = [
175        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
176        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
177        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
178        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
179        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
180        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
181        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
182        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
183    ];
184    while stream.read_chunks(&mut bufs[..]).await?.is_some() {}
185    debug!("finished reading {}", stream.id());
186    Ok(())
187}
188
189async fn respond(mut bytes: u64, mut stream: quinn::SendStream) -> Result<()> {
190    static DATA: [u8; 1024 * 1024] = [42; 1024 * 1024];
191
192    while bytes > 0 {
193        let chunk_len = bytes.min(DATA.len() as u64);
194        stream
195            .write_chunk(Bytes::from_static(&DATA[..chunk_len as usize]))
196            .await
197            .context("sending response")?;
198        bytes -= chunk_len;
199    }
200    debug!("finished responding on {}", stream.id());
201    Ok(())
202}
perf/server.rs

perf/
server.rs
