perf: don't allocate in UDP send & recv path

.github/workflows/polonius.diff

@@ -0,0 +1,50 @@
+diff --git a/neqo-http3/src/server.rs b/neqo-http3/src/server.rs
+index 9b73f2c7..d6cd39a2 100644
+--- a/neqo-http3/src/server.rs
++++ b/neqo-http3/src/server.rs
+@@ -120,12 +120,7 @@ impl Http3Server {
+         out: &'a mut Vec<u8>,
+     ) -> Output<&'a [u8]> {
+         qtrace!([self], "Process.");
+-        let mut output = self.server.process_into_buffer(
+-            dgram,
+-            now,
+-            // See .github/workflows/polonius.yml.
+-            unsafe { &mut *std::ptr::from_mut(out) },
+-        );
++        let mut output = self.server.process_into_buffer(dgram, now, out);
+
+         self.process_http3(now);
+
+diff --git a/neqo-transport/src/server.rs b/neqo-transport/src/server.rs
+index 2446f8d3..25ceb2b0 100644
+--- a/neqo-transport/src/server.rs
++++ b/neqo-transport/src/server.rs
+@@ -445,12 +445,7 @@ impl Server {
+         let mut callback = None;
+
+         for connection in &mut self.connections {
+-            match connection.borrow_mut().process_into_buffer(
+-                None,
+-                now,
+-                // See .github/workflows/polonius.yml.
+-                unsafe { &mut *std::ptr::from_mut(out) },
+-            ) {
++            match connection.borrow_mut().process_into_buffer(None, now, out) {
+                 Output::None => {}
+                 d @ Output::Datagram(_) => return d,
+                 Output::Callback(next) => match callback {
+@@ -491,12 +486,7 @@ impl Server {
+
+         #[allow(clippy::option_if_let_else)]
+         let output = if let Some(dgram) = dgram {
+-            self.process_input(
+-                dgram,
+-                now,
+-                // See .github/workflows/polonius.yml.
+-                unsafe { &mut *std::ptr::from_mut(out) },
+-            )
++            self.process_input(dgram, now, out)
+         } else {
+             Output::None
+         };

.github/workflows/polonius.yml

@@ -0,0 +1,61 @@
+# Rustc by default uses the NLL borrow-checker. There are multiple cases where
+# NLL is too restrictive. An upcoming improvement to Rust's borrow-checker,
+# adressing these shortcomings, is Polonius. It ships with Nightly behind a
+# flag.
+#
+# This workflow first removes the workarounds necessary to please NLL and then
+# runs with Polonius to ensure each workaround only fixes the false-positive of
+# NLL and doesn't mask an actually undefined behavior.
+#
+# See also:
+# - <https://blog.rust-lang.org/inside-rust/2023/10/06/polonius-update.html>
+# - <https://github.com/rust-lang/rust/issues/54663>
+
+name: Polonius
+on:
+  push:
+    branches: ["main"]
+    paths-ignore: ["*.md", "*.png", "*.svg", "LICENSE-*"]
+  pull_request:
+    branches: ["main"]
+    paths-ignore: ["*.md", "*.png", "*.svg", "LICENSE-*"]
+  merge_group:
+  workflow_dispatch:
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: 1
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref_name }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  polonius:
+    strategy:
+      fail-fast: false
+    defaults:
+      run:
+        shell: bash
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+      - uses: ./.github/actions/rust
+        with:
+          version: nightly
+          token: ${{ secrets.GITHUB_TOKEN }}
+
+      - id: nss-version
+        run: echo "minimum=$(cat neqo-crypto/min_version.txt)" >> "$GITHUB_OUTPUT"
+
+      - uses: ./.github/actions/nss
+        with:
+          minimum-version: ${{ steps.nss-version.outputs.minimum }}
+
+      - name: Apply patch, removing `unsafe` workarounds.
+        run: git apply .github/workflows/polonius.diff
+
+      - run: RUSTFLAGS="-Z polonius" cargo +nightly check

fuzz/fuzz_targets/client_initial.rs

@@ -24,11 +24,13 @@ fuzz_target!(|data: &[u8]| {
     let (aead, hp) = initial_aead_and_hp(d_cid, Role::Client);
     let (_, pn) = remove_header_protection(&hp, header, payload);
 
-    let mut payload_enc = Encoder::with_capacity(MIN_INITIAL_PACKET_SIZE);
+    let mut out = Vec::with_capacity(MIN_INITIAL_PACKET_SIZE);
+    let mut payload_enc = Encoder::new(&mut out);
     payload_enc.encode(data); // Add fuzzed data.
 
     // Make a new header with a 1 byte packet number length.
-    let mut header_enc = Encoder::new();
+    let mut out = Vec::new();
+    let mut header_enc = Encoder::new(&mut out);
     header_enc
         .encode_byte(0xc0) // Initial with 1 byte packet number.
         .encode_uint(4, Version::default().wire_version())
@@ -57,7 +59,7 @@ fuzz_target!(|data: &[u8]| {
         &mut ciphertext,
         (header_enc.len() - 1)..header_enc.len(),
     );
-    let fuzzed_ci = Datagram::new(ci.source(), ci.destination(), ci.tos(), ciphertext);
+    let fuzzed_ci = Datagram::new(ci.source(), ci.destination(), ci.tos(), ciphertext, None);
 
     let mut server = default_server();
     let _response = server.process(Some(&fuzzed_ci), now());

fuzz/fuzz_targets/server_initial.rs

@@ -30,11 +30,13 @@ fuzz_target!(|data: &[u8]| {
     let (aead, hp) = initial_aead_and_hp(d_cid, Role::Server);
     let (_, pn) = remove_header_protection(&hp, header, payload);
 
-    let mut payload_enc = Encoder::with_capacity(MIN_INITIAL_PACKET_SIZE);
+    let mut out = Vec::with_capacity(MIN_INITIAL_PACKET_SIZE);
+    let mut payload_enc = Encoder::new(&mut out);
     payload_enc.encode(data); // Add fuzzed data.
 
     // Make a new header with a 1 byte packet number length.
-    let mut header_enc = Encoder::new();
+    let mut out = Vec::new();
+    let mut header_enc = Encoder::new(&mut out);
     header_enc
         .encode_byte(0xc0) // Initial with 1 byte packet number.
         .encode_uint(4, Version::default().wire_version())
@@ -63,7 +65,7 @@ fuzz_target!(|data: &[u8]| {
         &mut ciphertext,
         (header_enc.len() - 1)..header_enc.len(),
     );
-    let fuzzed_si = Datagram::new(si.source(), si.destination(), si.tos(), ciphertext);
+    let fuzzed_si = Datagram::new(si.source(), si.destination(), si.tos(), ciphertext, None);
     let _response = client.process(Some(&fuzzed_si), now());
 });
 

neqo-bin/src/client/http09.rs

@@ -17,7 +17,7 @@ use std::{
     time::Instant,
 };
 
-use neqo_common::{event::Provider, qdebug, qinfo, qwarn, Datagram};
+use neqo_common::{event::Provider, qdebug, qinfo, qwarn, BorrowedDatagram};
 use neqo_crypto::{AuthenticationStatus, ResumptionToken};
 use neqo_transport::{
     CloseReason, Connection, ConnectionEvent, EmptyConnectionIdGenerator, Error, Output, State,
@@ -177,15 +177,13 @@ impl TryFrom<&State> for CloseState {
 }
 
 impl super::Client for Connection {
-    fn process_output(&mut self, now: Instant) -> Output {
-        self.process_output(now)
-    }
-
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>,
-    {
-        self.process_multiple_input(dgrams, now);
+    fn process_into_buffer<'a>(
+        &mut self,
+        input: Option<BorrowedDatagram>,
+        now: Instant,
+        out: &'a mut Vec<u8>,
+    ) -> Output<BorrowedDatagram<'a>> {
+        self.process_into_buffer(input, now, out)
     }
 
     fn close<S>(&mut self, now: Instant, app_error: neqo_transport::AppError, msg: S)

neqo-bin/src/client/http3.rs

@@ -18,7 +18,7 @@ use std::{
     time::Instant,
 };
 
-use neqo_common::{event::Provider, hex, qdebug, qinfo, qwarn, Datagram, Header};
+use neqo_common::{event::Provider, hex, qdebug, qinfo, qwarn, BorrowedDatagram, Header};
 use neqo_crypto::{AuthenticationStatus, ResumptionToken};
 use neqo_http3::{Error, Http3Client, Http3ClientEvent, Http3Parameters, Http3State, Priority};
 use neqo_transport::{
@@ -128,15 +128,13 @@ impl super::Client for Http3Client {
         self.state().try_into()
     }
 
-    fn process_output(&mut self, now: Instant) -> Output {
-        self.process_output(now)
-    }
-
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>,
-    {
-        self.process_multiple_input(dgrams, now);
+    fn process_into_buffer<'a>(
+        &mut self,
+        input: Option<BorrowedDatagram>,
+        now: Instant,
+        out: &'a mut Vec<u8>,
+    ) -> Output<BorrowedDatagram<'a>> {
+        self.process_into_buffer(input, now, out)
     }
 
     fn close<S>(&mut self, now: Instant, app_error: AppError, msg: S)

neqo-bin/src/client/mod.rs

@@ -23,7 +23,7 @@ use futures::{
     future::{select, Either},
     FutureExt, TryFutureExt,
 };
-use neqo_common::{qdebug, qerror, qinfo, qlog::NeqoQlog, qwarn, Datagram, Role};
+use neqo_common::{qdebug, qerror, qinfo, qlog::NeqoQlog, qwarn, BorrowedDatagram, Role};
 use neqo_crypto::{
     constants::{TLS_AES_128_GCM_SHA256, TLS_AES_256_GCM_SHA384, TLS_CHACHA20_POLY1305_SHA256},
     init, Cipher, ResumptionToken,
@@ -372,10 +372,12 @@ enum CloseState {
 
 /// Network client, e.g. [`neqo_transport::Connection`] or [`neqo_http3::Http3Client`].
 trait Client {
-    fn process_output(&mut self, now: Instant) -> Output;
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>;
+    fn process_into_buffer<'a>(
+        &mut self,
+        input: Option<BorrowedDatagram>,
+        now: Instant,
+        out: &'a mut Vec<u8>,
+    ) -> Output<BorrowedDatagram<'a>>;
     fn has_events(&self) -> bool;
     fn close<S>(&mut self, now: Instant, app_error: AppError, msg: S)
     where
@@ -391,13 +393,34 @@ struct Runner<'a, H: Handler> {
     handler: H,
     timeout: Option<Pin<Box<Sleep>>>,
     args: &'a Args,
+    recv_buf: Vec<u8>,
+    send_buf: Vec<u8>,
 }
 
 impl<'a, H: Handler> Runner<'a, H> {
+    fn new(
+        local_addr: SocketAddr,
+        socket: &'a mut crate::udp::Socket,
+        client: H::Client,
+        handler: H,
+        args: &'a Args,
+    ) -> Self {
+        Self {
+            local_addr,
+            socket,
+            client,
+            handler,
+            args,
+            timeout: None,
+            recv_buf: vec![0; neqo_udp::RECV_BUF_SIZE],
+            send_buf: Vec::new(),
+        }
+    }
+
     async fn run(mut self) -> Res<Option<ResumptionToken>> {
         loop {
             let handler_done = self.handler.handle(&mut self.client)?;
-            self.process_output().await?;
+            self.process(false).await?;
             if self.client.has_events() {
                 continue;
             }
@@ -418,7 +441,7 @@ impl<'a, H: Handler> Runner<'a, H> {
             }
 
             match ready(self.socket, self.timeout.as_mut()).await? {
-                Ready::Socket => self.process_multiple_input().await?,
+                Ready::Socket => self.process(true).await?,
                 Ready::Timeout => {
                     self.timeout = None;
                 }
@@ -432,37 +455,36 @@ impl<'a, H: Handler> Runner<'a, H> {
         Ok(self.handler.take_token())
     }
 
-    async fn process_output(&mut self) -> Result<(), io::Error> {
+    async fn process(&mut self, mut should_read: bool) -> Result<(), io::Error> {
         loop {
-            match self.client.process_output(Instant::now()) {
+            let dgram = should_read
+                .then(|| self.socket.recv(&self.local_addr, &mut self.recv_buf))
+                .transpose()?
+                .flatten();
+            should_read = dgram.is_some();
+
+            match self
+                .client
+                .process_into_buffer(dgram, Instant::now(), &mut self.send_buf)
+            {
                 Output::Datagram(dgram) => {
                     self.socket.writable().await?;
-                    self.socket.send(&dgram)?;
+                    self.socket.send(dgram)?;
+                    self.send_buf.clear();
+                    continue;
                 }
                 Output::Callback(new_timeout) => {
                     qdebug!("Setting timeout of {:?}", new_timeout);
                     self.timeout = Some(Box::pin(tokio::time::sleep(new_timeout)));
-                    break;
                 }
                 Output::None => {
                     qdebug!("Output::None");
-                    break;
                 }
             }
-        }
 
-        Ok(())
-    }
-
-    async fn process_multiple_input(&mut self) -> Res<()> {
-        loop {
-            let dgrams = self.socket.recv(&self.local_addr)?;
-            if dgrams.is_empty() {
+            if !should_read {
                 break;
             }
-            self.client
-                .process_multiple_input(dgrams.iter(), Instant::now());
-            self.process_output().await?;
         }
 
         Ok(())
@@ -572,32 +594,18 @@ pub async fn client(mut args: Args) -> Res<()> {
 
                 let handler = http09::Handler::new(to_request, &args);
 
-                Runner {
-                    args: &args,
-                    client,
-                    handler,
-                    local_addr: real_local,
-                    socket: &mut socket,
-                    timeout: None,
-                }
-                .run()
-                .await?
+                Runner::new(real_local, &mut socket, client, handler, &args)
+                    .run()
+                    .await?
             } else {
                 let client = http3::create_client(&args, real_local, remote_addr, &hostname, token)
                     .expect("failed to create client");
 
                 let handler = http3::Handler::new(to_request, &args);
 
-                Runner {
-                    args: &args,
-                    client,
-                    handler,
-                    local_addr: real_local,
-                    socket: &mut socket,
-                    timeout: None,
-                }
-                .run()
-                .await?
+                Runner::new(real_local, &mut socket, client, handler, &args)
+                    .run()
+                    .await?
             };
         }
     }