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 noq::{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 = QuicServerConfig::try_from(crypto)?;
71    let mut config = noq::ServerConfig::with_crypto(match opt.common.no_protection {
72        true => Arc::new(NoProtectionServerConfig::new(crypto)),
73        false => Arc::new(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 = noq::EndpointConfig::default();
80    endpoint_cfg.max_udp_payload_size(opt.common.max_udp_payload_size)?;
81
82    let endpoint = noq::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: noq::Incoming, opt: Arc<Opt>) -> Result<()> {
102    let connection = handshake.await.context("handshake failed")?;
103
104    debug!(
105        "{} connected",
106        connection
107            .path(noq::PathId::ZERO)
108            .expect("path open after connect")
109            .remote_address()
110            .expect("path is alive")
111    );
112    tokio::try_join!(
113        drive_uni(connection.clone()),
114        drive_bi(connection.clone()),
115        conn_stats(connection, opt)
116    )?;
117    Ok(())
118}
119
120async fn conn_stats(connection: noq::Connection, opt: Arc<Opt>) -> Result<()> {
121    if opt.common.conn_stats {
122        loop {
123            tokio::time::sleep(Duration::from_secs(2)).await;
124            println!("{:?}\n", connection.stats());
125        }
126    }
127
128    Ok(())
129}
130
131async fn drive_uni(connection: noq::Connection) -> Result<()> {
132    while let Ok(stream) = connection.accept_uni().await {
133        let connection = connection.clone();
134        tokio::spawn(async move {
135            if let Err(e) = handle_uni(connection, stream).await {
136                error!("request failed: {:#}", e);
137            }
138        });
139    }
140    Ok(())
141}
142
143async fn handle_uni(connection: noq::Connection, stream: noq::RecvStream) -> Result<()> {
144    let bytes = read_req(stream).await?;
145    let response = connection.open_uni().await?;
146    respond(bytes, response).await?;
147    Ok(())
148}
149
150async fn drive_bi(connection: noq::Connection) -> Result<()> {
151    while let Ok((send, recv)) = connection.accept_bi().await {
152        tokio::spawn(async move {
153            if let Err(e) = handle_bi(send, recv).await {
154                error!("request failed: {:#}", e);
155            }
156        });
157    }
158    Ok(())
159}
160
161async fn handle_bi(send: noq::SendStream, recv: noq::RecvStream) -> Result<()> {
162    let bytes = read_req(recv).await?;
163    respond(bytes, send).await?;
164    Ok(())
165}
166
167async fn read_req(mut stream: noq::RecvStream) -> Result<u64> {
168    let mut buf = [0; 8];
169    stream
170        .read_exact(&mut buf)
171        .await
172        .context("reading request")?;
173    let n = u64::from_be_bytes(buf);
174    debug!("got req for {} bytes on {}", n, stream.id());
175    drain_stream(stream).await?;
176    Ok(n)
177}
178
179async fn drain_stream(mut stream: noq::RecvStream) -> Result<()> {
180    #[rustfmt::skip]
181    let mut bufs = [
182        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
183        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
184        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
185        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
186        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
187        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
188        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
189        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
190    ];
191    while stream.read_many_chunks(&mut bufs[..]).await?.is_some() {}
192    debug!("finished reading {}", stream.id());
193    Ok(())
194}
195
196async fn respond(mut bytes: u64, mut stream: noq::SendStream) -> Result<()> {
197    static DATA: [u8; 1024 * 1024] = [42; 1024 * 1024];
198
199    while bytes > 0 {
200        let chunk_len = bytes.min(DATA.len() as u64);
201        stream
202            .write_chunk(Bytes::from_static(&DATA[..chunk_len as usize]))
203            .await
204            .context("sending response")?;
205        bytes -= chunk_len;
206    }
207    debug!("finished responding on {}", stream.id());
208    Ok(())
209}
perf/server.rs

perf/
server.rs
