Enable allocation of many packet bufs for client socket at a time

3 files changed, 16 insertions, 6 deletions

diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index cb178788ec..62139a1ab8 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -464,11 +464,16 @@ async fn client_socket_rx_task(
     client_tx: mpsc::Sender<Bytes>,
     return_addr_tx: broadcast::Sender<SocketAddr>,
 ) -> Result<()> {
-    let mut client_read_buf = BytesMut::with_capacity(100 * crate::PACKET_BUFFER_SIZE);
+    const TOTAL_BUFFER_CAPACITY: usize = 100 * crate::MAX_UDP_SIZE;
+
+    let mut client_read_buf = BytesMut::with_capacity(TOTAL_BUFFER_CAPACITY);
     let mut return_addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0);
 
     loop {
-        client_read_buf.reserve(crate::PACKET_BUFFER_SIZE);
+        if !client_read_buf.try_reclaim(crate::MAX_UDP_SIZE) {
+            // Allocate space for new packets
+            client_read_buf.reserve(TOTAL_BUFFER_CAPACITY);
+        }
 
         // this is the variable ID used to signify UDP payloads in HTTP datagrams.
         crate::HTTP_MASQUE_DATAGRAM_CONTEXT_ID.encode(&mut client_read_buf);
diff --git a/mullvad-masque-proxy/src/lib.rs b/mullvad-masque-proxy/src/lib.rs
index 37f8353acd..2c7b64f900 100644
--- a/mullvad-masque-proxy/src/lib.rs
+++ b/mullvad-masque-proxy/src/lib.rs
@@ -11,9 +11,9 @@ pub const MASQUE_WELL_KNOWN_PATH: &str = "/.well-known/masque/udp/";
 pub const HTTP_MASQUE_DATAGRAM_CONTEXT_ID: VarInt = VarInt::from_u32(0);
 pub const HTTP_MASQUE_FRAGMENTED_DATAGRAM_CONTEXT_ID: VarInt = VarInt::from_u32(1);
 
-/// Minimum size of buffer used to hold UDP packets.
+/// Maximum possible buffer size UDP packets, plus context ID.
 // 1 byte for size of HTTP_MASQUE_DATAGRAM_CONTEXT_ID
-const PACKET_BUFFER_SIZE: usize = (u16::MAX - UDP_HEADER_SIZE + 1) as usize;
+const MAX_UDP_SIZE: usize = (u16::MAX - UDP_HEADER_SIZE + 1) as usize;
 
 /// Maximum number of inflight packets, in both directions.
 const MAX_INFLIGHT_PACKETS: usize = 100;
diff --git a/mullvad-masque-proxy/src/server/mod.rs b/mullvad-masque-proxy/src/server/mod.rs
index 3788ca7c80..2232fbdae2 100644
--- a/mullvad-masque-proxy/src/server/mod.rs
+++ b/mullvad-masque-proxy/src/server/mod.rs
@@ -337,13 +337,18 @@ async fn proxy_rx_task(
     udp_socket: impl AsRef<UdpSocket>,
     send_tx: mpsc::Sender<Bytes>,
 ) {
+    const TOTAL_BUFFER_CAPACITY: usize = 100 * crate::MAX_UDP_SIZE;
+
     let stream_id_size = VarInt::from(stream_id).size() as u16;
     let udp_socket = udp_socket.as_ref();
-    let mut proxy_recv_buf = BytesMut::with_capacity(100 * crate::PACKET_BUFFER_SIZE);
+    let mut proxy_recv_buf = BytesMut::with_capacity(TOTAL_BUFFER_CAPACITY);
     let mut fragment_id = 0u16;
 
     loop {
-        proxy_recv_buf.reserve(crate::PACKET_BUFFER_SIZE);
+        if !proxy_recv_buf.try_reclaim(crate::MAX_UDP_SIZE) {
+            // Allocate space for new packets
+            proxy_recv_buf.reserve(TOTAL_BUFFER_CAPACITY);
+        }
         crate::HTTP_MASQUE_DATAGRAM_CONTEXT_ID.encode(&mut proxy_recv_buf);
 
         let (_n, sender_addr) = match udp_socket.recv_buf_from(&mut proxy_recv_buf).await {