Limit send flow control bug to window_size (#78)

Senders could set the available capacity greater than the current
`window_size`.  This caused a panic when the sender attempted
to send more than the receiver could accept.

src/codec/framed_read.rs

@@ -249,6 +249,7 @@ impl<T> Stream for FramedRead<T>
 
             trace!("poll; bytes={}B", bytes.len());
             if let Some(frame) = try!(self.decode_frame(bytes)) {
+                debug!("received; frame={:?}", frame);
                 return Ok(Async::Ready(Some(frame)));
             }
         }

src/proto/streams/prioritize.rs

@@ -246,7 +246,8 @@ impl<B, P> Prioritize<B, P>
         // Don't assign more than the window has available!
         let additional = cmp::min(
             total_requested - stream.send_flow.available(),
-            stream.send_flow.window_size());
+            // Can't assign more than what is available
+            stream.send_flow.window_size() - stream.send_flow.available());
 
         trace!("try_assign_capacity; requested={}; additional={}; buffered={}; window={}; conn={}",
                total_requested,
@@ -451,8 +452,6 @@ impl<B, P> Prioritize<B, P>
                         Frame::Data(mut frame) => {
                             // Get the amount of capacity remaining for stream's
                             // window.
-                            //
-                            // TODO: Is this the right thing to check?
                             let stream_capacity = stream.send_flow.available();
                             let sz = frame.payload().remaining();
 

tests/flow_control.rs

@@ -539,3 +539,83 @@ fn recv_window_update_on_stream_closed_by_data_frame() {
     let _ = h2.join(srv)
         .wait().unwrap();
 }
+
+#[test]
+fn reserved_capacity_assigned_in_multi_window_updates() {
+    let _ = ::env_logger::init();
+    let (io, srv) = mock::new();
+
+    let h2 = Client::handshake(io).unwrap()
+        .and_then(|mut h2| {
+            let request = Request::builder()
+                .method(Method::POST)
+                .uri("https://http2.akamai.com/")
+                .body(()).unwrap();
+
+            let mut stream = h2.request(request, false).unwrap();
+
+            // Consume the capacity
+            let payload = vec![0; frame::DEFAULT_INITIAL_WINDOW_SIZE as usize];
+            stream.send_data(payload.into(), false).unwrap();
+
+            // Reserve more data than we want
+            stream.reserve_capacity(10);
+
+            h2.drive(util::wait_for_capacity(stream, 5))
+        })
+        .and_then(|(h2, mut stream)| {
+            stream.send_data("hello".into(), false).unwrap();
+            stream.send_data("world".into(), true).unwrap();
+
+            h2.drive(GetResponse { stream: Some(stream) })
+        })
+        .and_then(|(h2, (response, _))| {
+            assert_eq!(response.status(), StatusCode::NO_CONTENT);
+
+            // Wait for the connection to close
+            h2.unwrap()
+        })
+        ;
+
+    let srv = srv.assert_client_handshake().unwrap()
+        .recv_settings()
+        .recv_frame(
+            frames::headers(1)
+                .request("POST", "https://http2.akamai.com/")
+        )
+        .recv_frame(frames::data(1, vec![0u8; 16_384]))
+        .recv_frame(frames::data(1, vec![0u8; 16_384]))
+        .recv_frame(frames::data(1, vec![0u8; 16_384]))
+        .recv_frame(frames::data(1, vec![0u8; 16_383]))
+        .idle_ms(100)
+        // Increase the connection window
+        .send_frame(
+            frames::window_update(0, 10))
+        // Incrementally increase the stream window
+        .send_frame(
+            frames::window_update(1, 4))
+        .idle_ms(50)
+        .send_frame(
+            frames::window_update(1, 1))
+        // Receive first chunk
+        .recv_frame(frames::data(1, "hello"))
+        .send_frame(
+            frames::window_update(1, 5))
+        // Receive second chunk
+        .recv_frame(
+            frames::data(1, "world").eos())
+        .send_frame(
+            frames::headers(1)
+                .response(204)
+                .eos()
+        )
+        /*
+        .recv_frame(frames::data(1, "hello").eos())
+        .send_frame(frames::window_update(1, 5))
+        */
+        .map(drop)
+        ;
+
+    let _ = h2.join(srv)
+        .wait().unwrap();
+}

tests/prioritization.rs

@@ -170,40 +170,6 @@ fn single_stream_send_extra_large_body_multi_frames_multi_buffer() {
     h2.wait().unwrap();
 }
 
-#[macro_use]
-extern crate futures;
-
-use futures::{Poll, Async};
-
-// TODO: Extract this out?
-struct WaitForCapacity {
-    stream: Option<client::Stream<Bytes>>,
-    target: usize,
-}
-
-impl WaitForCapacity {
-    fn stream(&mut self) -> &mut client::Stream<Bytes> {
-        self.stream.as_mut().unwrap()
-    }
-}
-
-impl Future for WaitForCapacity {
-    type Item = client::Stream<Bytes>;
-    type Error = ();
-
-    fn poll(&mut self) -> Poll<Self::Item, ()> {
-        let _ = try_ready!(self.stream().poll_capacity().map_err(|_| panic!()));
-
-        let act = self.stream().capacity();
-
-        if act >= self.target {
-            return Ok(self.stream.take().unwrap().into())
-        }
-
-        Ok(Async::NotReady)
-    }
-}
-
 #[test]
 fn send_data_receive_window_update() {
     let _ = ::env_logger::init();
@@ -225,10 +191,7 @@ fn send_data_receive_window_update() {
             stream.reserve_capacity(frame::DEFAULT_INITIAL_WINDOW_SIZE as usize);
 
             // Wait for capacity
-            h2.drive(WaitForCapacity {
+            h2.drive(util::wait_for_capacity(stream, frame::DEFAULT_INITIAL_WINDOW_SIZE as usize))
-                stream: Some(stream),
-                target: frame::DEFAULT_INITIAL_WINDOW_SIZE as usize,
-            })
         })
         .and_then(|(h2, mut stream)| {
             let payload = vec![0; frame::DEFAULT_INITIAL_WINDOW_SIZE as usize];

tests/support/src/lib.rs

@@ -6,6 +6,8 @@ pub extern crate http;
 
 #[macro_use]
 pub extern crate tokio_io;
+
+#[macro_use]
 pub extern crate futures;
 pub extern crate mock_io;
 pub extern crate env_logger;
@@ -19,6 +21,7 @@ pub mod raw;
 pub mod frames;
 pub mod prelude;
 pub mod mock;
+pub mod util;
 
 mod future_ext;
 

tests/support/src/mock.rs

@@ -5,6 +5,7 @@ use h2::frame::{self, Frame};
 
 use futures::{Async, Future, Stream, Poll};
 use futures::task::{self, Task};
+use futures::sync::oneshot;
 
 use tokio_io::{AsyncRead, AsyncWrite};
 use tokio_io::io::read_exact;
@@ -100,12 +101,18 @@ impl Handle {
     }
 
     /// Perform the H2 handshake
-    pub fn assert_client_handshake(mut self)
+    pub fn assert_client_handshake(self)
+        -> Box<Future<Item = (frame::Settings, Self), Error = h2::Error>>
+    {
+        self.assert_client_handshake_with_settings(frame::Settings::default())
+    }
+
+    /// Perform the H2 handshake
+    pub fn assert_client_handshake_with_settings(mut self, settings: frame::Settings)
         -> Box<Future<Item = (frame::Settings, Self), Error = h2::Error>>
     {
         // Send a settings frame
-        let frame = frame::Settings::default();
+        self.send(settings.into()).unwrap();
-        self.send(frame.into()).unwrap();
 
         let ret = self.read_preface().unwrap()
             .and_then(|me| me.into_future().unwrap())
@@ -336,6 +343,31 @@ pub trait HandleFutureExt {
         }
     }
 
+    fn idle_ms(self, ms: usize) -> Box<Future<Item=Handle, Error=Self::Error>>
+        where Self: Sized + 'static,
+              Self: Future<Item=Handle>,
+              Self::Error: fmt::Debug,
+    {
+        use std::thread;
+        use std::time::Duration;
+
+
+        Box::new(self.and_then(move |handle| {
+            // This is terrible... but oh well
+            let (tx, rx) = oneshot::channel();
+
+            thread::spawn(move || {
+                thread::sleep(Duration::from_millis(ms as u64));
+                tx.send(()).unwrap();
+            });
+
+            Idle {
+                handle: Some(handle),
+                timeout: rx,
+            }.map_err(|_| unreachable!())
+        }))
+    }
+
     fn close(self) -> Box<Future<Item=(), Error=()>>
         where Self: Future<Error = ()> + Sized + 'static,
     {
@@ -387,6 +419,29 @@ impl<T> Future for SendFrameFut<T>
     }
 }
 
+pub struct Idle {
+    handle: Option<Handle>,
+    timeout: oneshot::Receiver<()>,
+}
+
+impl Future for Idle {
+    type Item = Handle;
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        if self.timeout.poll().unwrap().is_ready() {
+            return Ok(self.handle.take().unwrap().into());
+        }
+
+        match self.handle.as_mut().unwrap().poll() {
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            res => {
+                panic!("Received unexpected frame on handle; frame={:?}", res);
+            }
+        }
+    }
+}
+
 impl<T> HandleFutureExt for T
     where T: Future + 'static,
 {

tests/support/src/prelude.rs

@@ -13,6 +13,9 @@ pub use super::mock::{self, HandleFutureExt};
 // Re-export frames helpers
 pub use super::frames;
 
+// Re-export utility mod
+pub use super::util;
+
 // Re-export some type defines
 pub use super::{Codec, SendFrame};
 

tests/support/src/util.rs

@@ -0,0 +1,44 @@
+use h2::client;
+
+use futures::{Poll, Async, Future};
+use bytes::Bytes;
+
+pub fn wait_for_capacity(stream: client::Stream<Bytes>,
+                         target: usize)
+    -> WaitForCapacity
+{
+    WaitForCapacity {
+        stream: Some(stream),
+        target: target,
+    }
+}
+
+pub struct WaitForCapacity {
+    stream: Option<client::Stream<Bytes>>,
+    target: usize,
+}
+
+impl WaitForCapacity {
+    fn stream(&mut self) -> &mut client::Stream<Bytes> {
+        self.stream.as_mut().unwrap()
+    }
+}
+
+impl Future for WaitForCapacity {
+    type Item = client::Stream<Bytes>;
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<Self::Item, ()> {
+        let _ = try_ready!(self.stream().poll_capacity().map_err(|_| panic!()));
+
+        let act = self.stream().capacity();
+
+        println!("CAP={:?}", act);
+
+        if act >= self.target {
+            return Ok(self.stream.take().unwrap().into())
+        }
+
+        Ok(Async::NotReady)
+    }
+}