Musitdev/prevent untrusted memory allocation

network/netcore/src/transport/proxy_protocol.rs

@@ -82,45 +82,63 @@ pub async fn read_header<T: AsyncRead + std::marker::Unpin>(
     let address_size: [u8; 2] = header[14..16].try_into().unwrap();
     let address_size = u16::from_be_bytes(address_size);
 
-    let mut address_bytes: Vec<u8> = vec![0; address_size as usize];
-    stream.read_exact(&mut address_bytes).await?;
+    // Determine how many bytes of the declared address region we actually need to
+    // parse. This is a compile-time constant per address family — never the raw
+    // wire value — so no allocation is ever proportional to untrusted input.
+    let fixed_parse_size: u16 = match family_and_protocol {
+        TCP_IPV4 if address_size >= IPV4_SIZE => IPV4_SIZE,
+        TCP_IPV6 if address_size >= IPV6_SIZE => IPV6_SIZE,
+        _ => 0,
+    };
+
+    // Read only the fixed portion we need into a stack buffer sized to the
+    // largest known address type. The remaining declared bytes are drained below.
+    let mut addr_buf = [0u8; IPV6_SIZE as usize];
+    if fixed_parse_size > 0 {
+        stream
+            .read_exact(&mut addr_buf[..fixed_parse_size as usize])
+            .await?;
+    }
+
+    // Drain whatever the sender declared beyond what we parsed, in bounded chunks.
+    // This single drain covers every branch: LOCAL/UDP/UNIX (all bytes), IPv4/IPv6
+    // with exact size (zero bytes), IPv4/IPv6 with extra padding, and error cases.
+    let remaining = address_size.saturating_sub(fixed_parse_size);
+    let mut scratch = [0u8; 256];
+    let mut left = remaining as usize;
+    while left > 0 {
+        let take = left.min(scratch.len());
+        stream.read_exact(&mut scratch[..take]).await?;
+        left = left.saturating_sub(take);
+    }
 
+    // Construct the result purely from the fixed buffer — the stream is fully consumed.
     let source_address = match family_and_protocol {
         // TODO: Support UDP in the future
         LOCAL_PROTOCOL | UDP_IPV4 | UDP_IPV6 | TCP_UNIX | UDP_UNIX => {
-            // UNSPEC, UDP, and UNIX Steam/datagram
+            // UNSPEC, UDP, and UNIX stream/datagram
             // Accept connection but ignore address info as per spec
             original_addr.clone()
         },
-        TCP_IPV4 => {
-            // This is not mentioned in the spec, but if it doesn't match we might not read correctly
-            if address_size < IPV4_SIZE {
-                return Err(io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "ProxyProtocol: Header size doesn't match expected address type",
-                ));
-            }
-
-            let src_addr = u32::from_be_bytes(address_bytes[0..4].try_into().unwrap());
-            let src_port = u16::from_be_bytes(address_bytes[8..10].try_into().unwrap());
+        TCP_IPV4 if address_size >= IPV4_SIZE => {
+            let src_addr = u32::from_be_bytes(addr_buf[0..4].try_into().unwrap());
+            let src_port = u16::from_be_bytes(addr_buf[8..10].try_into().unwrap());
             let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(src_addr)), src_port);
             NetworkAddress::from(socket_addr)
         },
-        TCP_IPV6 => {
-            // This is not mentioned in the spec, but if it doesn't match we might not read correctly
-            if address_size < IPV6_SIZE {
-                return Err(io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "ProxyProtocol: Header size doesn't match expected address type",
-                ));
-            }
-
-            let src_addr = u128::from_be_bytes(address_bytes[0..16].try_into().unwrap());
-            let src_port = u16::from_be_bytes(address_bytes[32..34].try_into().unwrap());
-
+        TCP_IPV6 if address_size >= IPV6_SIZE => {
+            let src_addr = u128::from_be_bytes(addr_buf[0..16].try_into().unwrap());
+            let src_port = u16::from_be_bytes(addr_buf[32..34].try_into().unwrap());
             let socket_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::from(src_addr)), src_port);
             NetworkAddress::from(socket_addr)
         },
+        TCP_IPV4 | TCP_IPV6 => {
+            // This is not mentioned in the spec, but if it doesn't match we might not read correctly
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "ProxyProtocol: Header size doesn't match expected address type",
+            ));
+        },
         _ => {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidInput,
@@ -136,8 +154,16 @@ pub async fn read_header<T: AsyncRead + std::marker::Unpin>(
 mod test {
     use super::*;
     use aptos_memsocket::MemorySocket;
-    use futures::{executor::block_on, future::join, io::AsyncWriteExt};
-    use std::net::ToSocketAddrs;
+    use futures::{
+        executor::block_on,
+        future::join,
+        io::{AsyncWriteExt, Cursor},
+    };
+    use std::{
+        net::ToSocketAddrs,
+        pin::Pin,
+        task::{Context, Poll},
+    };
 
     const TEST_DATA: &[u8; 4] = &[0xDE, 0xAD, 0xBE, 0xEF];
     const IPV4_ADDR_1: &[u8; 4] = &[0x00, 0x00, 0x00, 0x01];
@@ -154,6 +180,48 @@ mod test {
     const IPV4_ADDR_SIZE: &[u8; 2] = &[0x00, IPV4_SIZE as u8];
     const IPV6_ADDR_SIZE: &[u8; 2] = &[0x00, IPV6_SIZE as u8];
 
+    struct GuardedReadStream {
+        inner: Cursor<Vec<u8>>,
+        max_allowed_read_len: usize,
+        max_observed_read_len: usize,
+    }
+
+    impl GuardedReadStream {
+        fn new(data: Vec<u8>, max_allowed_read_len: usize) -> Self {
+            Self {
+                inner: Cursor::new(data),
+                max_allowed_read_len,
+                max_observed_read_len: 0,
+            }
+        }
+
+        fn max_observed_read_len(&self) -> usize {
+            self.max_observed_read_len
+        }
+    }
+
+    impl AsyncRead for GuardedReadStream {
+        fn poll_read(
+            mut self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &mut [u8],
+        ) -> Poll<io::Result<usize>> {
+            self.max_observed_read_len = self.max_observed_read_len.max(buf.len());
+
+            if buf.len() > self.max_allowed_read_len {
+                return Poll::Ready(Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!(
+                        "oversized read request from untrusted length field: {}",
+                        buf.len()
+                    ),
+                )));
+            }
+
+            Pin::new(&mut self.inner).poll_read(cx, buf)
+        }
+    }
+
     async fn send_v4(sender: &mut MemorySocket) -> io::Result<()> {
         sender.write_all(&PPV2_SIGNATURE).await?; // V2 signature
         sender.write_all(&[PPV2_PROXY]).await?; // Version 2, Proxy
@@ -246,10 +314,13 @@ mod test {
 
     #[test]
     fn test_local_proxy_protocol() {
-        let address_bytes: [&[u8; 1]; 5] =
-            [&[LOCAL_PROTOCOL], &[UDP_IPV4], &[UDP_IPV6], &[TCP_UNIX], &[
-                UDP_UNIX,
-            ]];
+        let address_bytes: [&[u8; 1]; 5] = [
+            &[LOCAL_PROTOCOL],
+            &[UDP_IPV4],
+            &[UDP_IPV6],
+            &[TCP_UNIX],
+            &[UDP_UNIX],
+        ];
 
         address_bytes
             .iter()
@@ -292,6 +363,40 @@ mod test {
         block_on(check_data);
     }
 
+    #[test]
+    fn test_large_untrusted_length_uses_bounded_read_chunks() {
+        let original_addr = NetworkAddress::mock();
+
+        // Build a valid PROXY v2 header: LOCAL_PROTOCOL with address_size = u16::MAX.
+        // No address bytes follow — the stream ends after the 16-byte fixed header.
+        let mut input = Vec::new();
+        input.extend_from_slice(&PPV2_SIGNATURE);
+        input.push(PPV2_PROXY);
+        input.push(LOCAL_PROTOCOL);
+        input.extend_from_slice(&u16::MAX.to_be_bytes());
+
+        let max_allowed_read_len = 1024;
+        let mut stream = GuardedReadStream::new(input, max_allowed_read_len);
+
+        // The stream is truncated, so parsing must fail with UnexpectedEof.
+        // If the parser allocates a 65535-byte buffer and tries to read it at once,
+        // the guard fires with InvalidData instead — causing this assertion to fail.
+        let err = block_on(read_header(&original_addr, &mut stream))
+            .expect_err("truncated payload should fail");
+        assert_eq!(
+            err.kind(),
+            io::ErrorKind::UnexpectedEof,
+            "expected UnexpectedEof but got {:?} — parser may be making oversized read requests",
+            err
+        );
+        assert!(
+            stream.max_observed_read_len() <= max_allowed_read_len,
+            "parser requested a read of {} bytes, exceeding the {} byte bound",
+            stream.max_observed_read_len(),
+            max_allowed_read_len,
+        );
+    }
+
     #[test]
     fn test_error_handling() {
         // Bad Header