From e703d9d64310245616aa4873be61f4e1a3f2128a Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Fri, 26 Mar 2021 16:55:14 -0700
Subject: [PATCH] Fix missing flow control credit when reading a stream's tail

---
 quinn-proto/src/connection/streams/recv.rs  |  5 ++-
 quinn-proto/src/connection/streams/state.rs | 41 +++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/quinn-proto/src/connection/streams/recv.rs b/quinn-proto/src/connection/streams/recv.rs
index 0ca830094..804e82be8 100644
--- a/quinn-proto/src/connection/streams/recv.rs
+++ b/quinn-proto/src/connection/streams/recv.rs
@@ -314,7 +314,10 @@ impl<'a> Chunks<'a> {
             }
             ChunksState::Finished => {
                 debug_assert!(!drop);
-                ShouldTransmit(true)
+                // MAX_DATA may need to be issued, but MAX_STREAM_DATA is pointless
+                let max_data = self.streams.add_read_credits(self.read);
+                self.pending.max_data |= max_data.0;
+                max_data
             }
             ChunksState::Error(_, should_transmit) => {
                 debug_assert!(!drop);
diff --git a/quinn-proto/src/connection/streams/state.rs b/quinn-proto/src/connection/streams/state.rs
index 8bfbe7a72..4b3bebd2a 100644
--- a/quinn-proto/src/connection/streams/state.rs
+++ b/quinn-proto/src/connection/streams/state.rs
@@ -777,6 +777,47 @@ mod tests {
         )
     }
 
+    #[test]
+    fn trivial_flow_control() {
+        let mut client = make(Side::Client);
+        let id = StreamId::new(Side::Server, Dir::Uni, 0);
+        let initial_max = client.local_max_data;
+        const MESSAGE_SIZE: usize = 2048;
+        assert_eq!(
+            client
+                .received(
+                    frame::Stream {
+                        id,
+                        offset: 0,
+                        fin: true,
+                        data: Bytes::from_static(&[0; MESSAGE_SIZE]),
+                    },
+                    2048
+                )
+                .unwrap(),
+            ShouldTransmit(false)
+        );
+        assert_eq!(client.data_recvd, 2048);
+        assert_eq!(client.local_max_data - initial_max, 0);
+
+        let mut pending = Retransmits::default();
+        let mut recv = RecvStream {
+            id,
+            state: &mut client,
+            pending: &mut pending,
+        };
+
+        let mut chunks = recv.read(true).unwrap();
+        assert_eq!(
+            chunks.next(MESSAGE_SIZE).unwrap().unwrap().bytes.len(),
+            MESSAGE_SIZE
+        );
+        assert!(chunks.next(0).unwrap().is_none());
+        let _ = chunks.finalize();
+        assert!(pending.max_uni_stream_id);
+        assert_eq!(client.local_max_data - initial_max, MESSAGE_SIZE as u64);
+    }
+
     #[test]
     fn reset_flow_control() {
         let mut client = make(Side::Client);