fix: flush data buffered by send_data before the next write and before finish

h3-quinn/src/lib.rs

@@ -531,10 +531,10 @@ where
     }
 
     #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
-    fn poll_finish(
-        &mut self,
-        _cx: &mut task::Context<'_>,
-    ) -> Poll<Result<(), StreamErrorIncoming>> {
+    fn poll_finish(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
+        // Flush whatever `send_data` buffered before sending FIN, otherwise the tail of the
+        // last frame is dropped and the peer sees a truncated stream.
+        ready!(self.poll_ready(cx))?;
         Poll::Ready(
             self.stream
                 .finish()

h3/src/stream.rs

@@ -34,6 +34,12 @@ where
     D: Into<WriteBuf<B>>,
     B: Buf,
 {
+    // This future is not cancel-safe: if a previous `write` was dropped between `send_data` and
+    // the completion of `poll_ready`, the transport still holds the rest of that buffer. Flush it
+    // first, so that `send_data` is never called while the stream is not ready. Transports treat
+    // that as a misuse and escalate it to closing the whole connection. When nothing is buffered,
+    // `poll_ready` is ready immediately.
+    future::poll_fn(|cx| stream.poll_ready(cx)).await?;
     stream.send_data(data)?;
     future::poll_fn(|cx| stream.poll_ready(cx)).await?;
 

h3/src/tests/connection.rs

@@ -982,3 +982,41 @@ where
         .unwrap();
     stream.finish().await.unwrap();
 }
+
+/// `poll_finish` must flush whatever `send_data` handed to the transport before it sends `FIN`.
+/// Otherwise the tail of the last frame is silently dropped and the peer sees a truncated stream.
+#[tokio::test]
+async fn h3_quinn_finish_flushes_buffered_data() {
+    init_tracing();
+    let mut pair = Pair::default();
+    let server = pair.server_inner();
+
+    let server_fut = async {
+        let conn = server.accept().await.unwrap().await.unwrap();
+        let (_send, mut recv) = conn.accept_bi().await.expect("accept_bi");
+        recv.read_to_end(usize::MAX).await.expect("read_to_end")
+    };
+
+    let client_fut = async {
+        let mut client = pair.client().await;
+        let mut stream = future::poll_fn(|cx| {
+            <h3_quinn::Connection as quic::OpenStreams<Bytes>>::poll_open_bidi(&mut client, cx)
+        })
+        .await
+        .expect("open_bidi");
+        // `send_data` only hands the frame to the transport; nothing is on the wire yet.
+        stream
+            .send_data(Frame::Data(Bytes::from_static(b"payload")))
+            .expect("send_data");
+        // Finish right away, without a `poll_ready` in between.
+        future::poll_fn(|cx| stream.poll_finish(cx))
+            .await
+            .expect("finish");
+        // Keep the connection alive until the server has read everything.
+        (client, stream)
+    };
+
+    let (received, _client) = tokio::join!(server_fut, client_fut);
+    // DATA frame: type 0x00, length 7, then the payload.
+    assert_eq!(received, b"\x00\x07payload");
+}

h3/src/tests/request.rs

@@ -1,4 +1,4 @@
-use std::{hint::black_box, time::Duration};
+use std::{future::Future, hint::black_box, pin::pin, task::Poll, time::Duration};
 
 use assert_matches::assert_matches;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -1588,3 +1588,77 @@ where
         assert_matches!(server_result_stream, Ok(()));
     }
 }
+
+/// A `send_data` future that is dropped before it completes (a `timeout` or `select!` cancelling
+/// it) leaves the rest of the frame buffered in the transport. The next write on that stream, and
+/// `finish`, must flush it instead of treating the stream as misused and closing the connection.
+#[tokio::test]
+async fn finish_after_cancelled_send_data_keeps_connection_open() {
+    init_tracing();
+    let mut pair = Pair::default();
+    let mut server = pair.server();
+
+    // Larger than the peer's initial stream receive window, so one poll cannot write all of it.
+    const BODY_LEN: usize = 4 * 1024 * 1024;
+
+    let client_fut = async {
+        let (mut driver, mut client) = client::new(pair.client().await).await.expect("client init");
+        let drive_fut = async { future::poll_fn(|cx| driver.poll_close(cx)).await };
+        let req_fut = async move {
+            let mut request_stream = client
+                .send_request(Request::get("http://localhost/big").body(()).unwrap())
+                .await
+                .expect("request");
+
+            let response = request_stream.recv_response().await.expect("recv response");
+            assert_eq!(response.status(), StatusCode::OK);
+
+            let mut received = 0;
+            while let Some(chunk) = request_stream.recv_data().await.expect("recv data") {
+                received += chunk.remaining();
+            }
+            assert_eq!(received, BODY_LEN);
+        };
+        tokio::join!(req_fut, drive_fut)
+    };
+
+    let server_fut = async {
+        let conn = server.next().await;
+        let mut incoming_req = server::Connection::new(conn).await.unwrap();
+
+        let (_request, mut request_stream) = get_stream_blocking(&mut incoming_req)
+            .await
+            .expect("accept");
+        request_stream
+            .send_response(
+                Response::builder()
+                    .status(200)
+                    .body(())
+                    .expect("build response"),
+            )
+            .await
+            .expect("send_response");
+
+        // Poll `send_data` exactly once and then drop it, like a cancelled `timeout` would.
+        {
+            let mut send = pin!(request_stream.send_data(Bytes::from(vec![0u8; BODY_LEN])));
+            future::poll_fn(|cx| match send.as_mut().poll(cx) {
+                Poll::Pending => Poll::Ready(()),
+                Poll::Ready(res) => panic!("send_data completed in a single poll: {res:?}"),
+            })
+            .await;
+        }
+
+        // The rest of the body is still buffered. This is the first stream on the connection, so
+        // `finish` also writes a grease frame first; both must flush rather than fail.
+        request_stream.finish().await.expect("finish");
+
+        assert_matches!(
+            incoming_req.accept().await.err().unwrap(),
+            ConnectionError::Remote(ConnectionErrorIncoming::ApplicationClose{error_code: code, ..})
+            if code == Code::H3_NO_ERROR.value()
+        );
+    };
+
+    tokio::join!(server_fut, client_fut);
+}