From 17bebe719ad209c4f58a3229088b64a17193d034 Mon Sep 17 00:00:00 2001
From: Oliver Gould <ver@buoyant.io>
Date: Fri, 8 Sep 2017 12:24:06 -0700
Subject: [PATCH] Increment stream window capacity when sending window update
 (#67)

---
 src/proto/streams/recv.rs |  5 ++++-
 tests/flow_control.rs     | 21 +++++++++------------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/proto/streams/recv.rs b/src/proto/streams/recv.rs
index 252a182..d814d22 100644
--- a/src/proto/streams/recv.rs
+++ b/src/proto/streams/recv.rs
@@ -525,7 +525,7 @@ impl<B, P> Recv<B, P>
             try_ready!(dst.poll_ready());
 
             // Get the next stream
-            let stream = match self.pending_window_updates.pop(store) {
+            let mut stream = match self.pending_window_updates.pop(store) {
                 Some(stream) => stream,
                 None => return Ok(().into()),
             };
@@ -543,6 +543,9 @@ impl<B, P> Recv<B, P>
 
                 // Buffer it
                 dst.buffer(frame.into()).ok().expect("invalid WINDOW_UPDATE frame");
+
+                // Update flow control
+                stream.recv_flow.inc_window(incr).ok().expect("unexpected flow control state");
             }
         }
     }
diff --git a/tests/flow_control.rs b/tests/flow_control.rs
index c73e0c2..be17b1b 100644
--- a/tests/flow_control.rs
+++ b/tests/flow_control.rs
@@ -54,7 +54,7 @@ fn send_data_without_requesting_capacity() {
 fn release_capacity_sends_window_update() {
     let _ = ::env_logger::init();
 
-    let payload = vec![0u8; 65_535];
+    let payload = vec![0u8; 16_384];
 
     let (io, srv) = mock::new();
 
@@ -69,19 +69,16 @@ fn release_capacity_sends_window_update() {
             frames::headers(1)
                 .response(200)
         )
-        .send_frame(frames::data(1, &payload[0..16_384]))
-        .send_frame(frames::data(1, &payload[16_384..16_384 * 2]))
-        .send_frame(frames::data(1, &payload[16_384 * 2..16_384 * 3]))
+        .send_frame(frames::data(1, &payload[..]))
+        .send_frame(frames::data(1, &payload[..]))
+        .send_frame(frames::data(1, &payload[..]))
         .recv_frame(
             frames::window_update(0, 32_768)
         )
         .recv_frame(
             frames::window_update(1, 32_768)
         )
-        .send_frame(
-            frames::data(1, &payload[16_384 * 3..])
-                .eos()
-        )
+        .send_frame(frames::data(1, &payload[..]).eos())
         // gotta end the connection
         .map(drop);
 
@@ -103,21 +100,21 @@ fn release_capacity_sends_window_update() {
 
                 // read some body to use up window size to below half
                 .and_then(|(buf, body)| {
-                    assert_eq!(buf.unwrap().len(), 16_384);
+                    assert_eq!(buf.unwrap().len(), payload.len());
                     body.into_future().unwrap()
                 })
                 .and_then(|(buf, body)| {
-                    assert_eq!(buf.unwrap().len(), 16_384);
+                    assert_eq!(buf.unwrap().len(), payload.len());
                     body.into_future().unwrap()
                 })
                 .and_then(|(buf, mut body)| {
                     let buf = buf.unwrap();
-                    assert_eq!(buf.len(), 16_384);
+                    assert_eq!(buf.len(), payload.len());
                     body.release_capacity(buf.len() * 2).unwrap();
                     body.into_future().unwrap()
                 })
                 .and_then(|(buf, _)| {
-                    assert_eq!(buf.unwrap().len(), 16_383);
+                    assert_eq!(buf.unwrap().len(), payload.len());
                     Ok(())
                 });
             h2.unwrap().join(req)