fix(s2n-quic-dc): correctly process IPv6 on recv path (#2626)

camshaft · web-flow · commit 6c9ec17215b6 · 2025-05-06T10:08:38.000-06:00
diff --git a/dc/s2n-quic-dc/src/stream/application.rs b/dc/s2n-quic-dc/src/stream/application.rs
@@ -150,6 +150,18 @@ where
         self.read.protocol()
     }
 
+    #[inline]
+    pub fn set_read_mode(&mut self, read_mode: recv::ReadMode) -> &mut Self {
+        self.read.set_read_mode(read_mode);
+        self
+    }
+
+    #[inline]
+    pub fn set_ack_mode(&mut self, ack_mode: recv::AckMode) -> &mut Self {
+        self.read.set_ack_mode(ack_mode);
+        self
+    }
+
     #[inline]
     pub async fn write_from(
         &mut self,
diff --git a/dc/s2n-quic-dc/src/stream/client/rpc.rs b/dc/s2n-quic-dc/src/stream/client/rpc.rs
@@ -1,7 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{event, stream::application::Stream};
+use crate::{
+    event,
+    stream::{application::Stream, recv::application::ReadMode},
+};
 use s2n_quic_core::buffer::{self, writer::Storage};
 use std::{
     future::{poll_fn, Future},
@@ -20,6 +23,9 @@ where
 {
     let (mut reader, mut writer) = stream.into_split();
 
+    // prefer draining all of the packets before sending an ACK
+    reader.set_read_mode(ReadMode::UntilFull);
+
     // TODO if the request is large enough, should we spawn a task for it?
     let mut writer = async move {
         while !request.buffer_is_empty() {
diff --git a/dc/s2n-quic-dc/src/stream/recv/application.rs b/dc/s2n-quic-dc/src/stream/recv/application.rs
@@ -135,6 +135,18 @@ where
         self.0.sockets.protocol()
     }
 
+    #[inline]
+    pub fn set_read_mode(&mut self, read_mode: ReadMode) -> &mut Self {
+        self.0.read_mode = read_mode;
+        self
+    }
+
+    #[inline]
+    pub fn set_ack_mode(&mut self, ack_mode: AckMode) -> &mut Self {
+        self.0.ack_mode = ack_mode;
+        self
+    }
+
     #[inline]
     pub async fn read_into<S>(&mut self, out_buf: &mut S) -> io::Result<usize>
     where
diff --git a/dc/s2n-quic-dc/src/stream/recv/application/builder.rs b/dc/s2n-quic-dc/src/stream/recv/application/builder.rs
@@ -14,6 +14,8 @@ use crate::{
 use core::mem::ManuallyDrop;
 use s2n_quic_core::endpoint;
 
+use super::ReadMode;
+
 pub struct Builder<Sub> {
     endpoint: endpoint::Type,
     runtime: runtime::ArcHandle<Sub>,
@@ -44,7 +46,14 @@ where
         };
         let gso = shared.gso.clone();
         let send_buffer = msg::send::Message::new(remote_addr, gso);
-        let read_mode = Default::default();
+        // If the transport is reliable then it's handling ACKs. Otherwise, the application is sending
+        // ACKs so we want to do a little more compute per `read` call, if the application buffer allows
+        // for it.
+        let read_mode = if is_reliable {
+            ReadMode::Once
+        } else {
+            ReadMode::UntilFull
+        };
         let ack_mode = Default::default();
         let local_state = match endpoint {
             // reliable transports on the client need to read at least one packet in order to
diff --git a/dc/s2n-quic-dc/src/stream/socket/fd/udp.rs b/dc/s2n-quic-dc/src/stream/socket/fd/udp.rs
@@ -115,7 +115,10 @@ fn recv_msghdr(
 
     let len = exec(&mut msg)?;
 
+    // make sure the CMSG has the correct length
     cmsg.with_msg(&msg);
+    // make sure the addr has the correct length
+    addr.update_with_msg(&msg);
 
     Ok(len)
 }

Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,10 @@ fn recv_msghdr(`
`115`	`115`
`116`	`116`	`let len = exec(&mut msg)?;`
`117`	`117`
	`118`	`+ // make sure the CMSG has the correct length`
`118`	`119`	`cmsg.with_msg(&msg);`
	`120`	`+ // make sure the addr has the correct length`
	`121`	`+ addr.update_with_msg(&msg);`
`119`	`122`
`120`	`123`	`Ok(len)`
`121`	`124`	`}`