Merge branch 'use-split_first_chunk'

.github/workflows/ci.yml

@@ -68,7 +68,8 @@ jobs:
             rust: nightly
           - os: ubuntu-latest
             # MSRV. Not considered breaking when this has to be bumped.
-            rust: 1.70.0
+            # But update `rust-version` in `Cargo.toml`
+            rust: 1.77.0
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v2

Cargo.toml

@@ -6,6 +6,7 @@ license = "MIT OR Apache-2.0"
 description = "Tunnel UDP traffic inside a TCP stream. Each datagram is prefixed with a 16 bit unsigned integer containing the length"
 repository = "https://github.com/mullvad/udp-over-tcp"
 edition = "2021"
+rust-version = "1.77.0"
 publish = false
 
 [[bin]]

src/forward_traffic.rs

@@ -3,7 +3,7 @@ use err_context::BoxedErrorExt as _;
 use err_context::ResultExt as _;
 use futures::future::select;
 use futures::pin_mut;
-use std::convert::{Infallible, TryFrom};
+use std::convert::Infallible;
 use std::future::Future;
 use std::io;
 use std::mem;
@@ -101,36 +101,40 @@ async fn maybe_timeout<F: Future>(
 /// Forward all complete datagrams in `buffer` to `udp_out`.
 /// Returns the number of processed bytes.
 async fn forward_datagrams_in_buffer(udp_out: &UdpSocket, buffer: &[u8]) -> io::Result<usize> {
-    let mut header_start = 0;
+    let mut unprocessed_buffer = buffer;
     loop {
-        let header_end = header_start + HEADER_LEN;
-        // "parse" the header
-        let header = match buffer.get(header_start..header_end) {
-            Some(header) => <[u8; HEADER_LEN]>::try_from(header).unwrap(),
-            // Buffer does not contain entire header for next datagram
-            None => break Ok(header_start),
-        };
-        let datagram_len = usize::from(u16::from_be_bytes(header));
-        let datagram_start = header_end;
-        let datagram_end = datagram_start + datagram_len;
-
-        let datagram_data = match buffer.get(datagram_start..datagram_end) {
-            Some(datagram_data) => datagram_data,
+        let Some((datagram_data, tail)) = split_first_datagram(unprocessed_buffer) else {
             // The buffer does not contain the entire datagram
-            None => break Ok(header_start),
+            break Ok(buffer.len() - unprocessed_buffer.len());
         };
 
         let udp_write_len = udp_out.send(datagram_data).await?;
         assert_eq!(
-            udp_write_len, datagram_len,
+            udp_write_len,
+            datagram_data.len(),
             "Did not send entire UDP datagram"
         );
-        log::trace!("Forwarded {} byte TCP->UDP", datagram_len);
+        log::trace!("Forwarded {} bytes TCP->UDP", datagram_data.len());
 
-        header_start = datagram_end;
+        unprocessed_buffer = tail;
     }
 }
 
+/// Parses the header at the beginning of the `buffer` and if it contains a full
+/// `udp-to-tcp` datagram it splits the buffer and returns the datagram data and
+/// buffer tail as two separate slices: `(datagram_data, tail)`
+fn split_first_datagram(buffer: &[u8]) -> Option<(&[u8], &[u8])> {
+    let (header, tail) = buffer.split_first_chunk::<HEADER_LEN>()?;
+    let datagram_len = usize::from(u16::from_be_bytes(*header));
+
+    // TODO: These two get calls (and thus double bounds check) can be replaced with
+    // `split_at_checked` when stabilized: https://github.com/rust-lang/rust/issues/119128
+    let datagram_data = tail.get(..datagram_len)?;
+    let tail = tail.get(datagram_len..)?;
+
+    Some((datagram_data, tail))
+}
+
 /// Reads datagrams from `udp_in` and writes them (with the 16 bit header containing the length)
 /// to `tcp_out` indefinitely, or until an IO error happens on either socket.
 async fn process_udp2tcp(