From 5a8de5dc7d8fe3179d10158ea9a78ea3e41b81e3 Mon Sep 17 00:00:00 2001
From: Kevin Butler <haqkrs@gmail.com>
Date: Fri, 8 Apr 2016 00:17:11 +0100
Subject: [PATCH] fix(response): prevent keep-alive if the Request had a body

This puts a bandaid on #326. A more complete fix will require
some shuffling of the api exposed by Request so that we can
monitor for when the body has been fully read (or read up to
a defined limit).
---
 src/server.rs              | 17 +++++++++++---
 tests/examples/chaining.rs | 45 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/src/server.rs b/src/server.rs
index b7e15e2352..dacc5eb70f 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,8 +1,10 @@
+
 use std::net::ToSocketAddrs;
 use std::sync::{Arc, RwLock};
 use std::collections::HashMap;
 use std::time::Duration;
 use hyper::Result as HttpResult;
+use hyper::header::{ContentLength, Connection};
 use hyper::server::{Request, Response, Handler, Listening};
 use hyper::server::Server as HyperServer;
 
@@ -21,12 +23,21 @@ struct ArcServer<D>(Arc<Server<D>>);
 
 impl<D: Sync + Send + 'static> Handler for ArcServer<D> {
     fn handle<'a, 'k>(&'a self, req: Request<'a, 'k>, res: Response<'a>) {
+        let request_has_body = match req.headers.get::<ContentLength>() {
+            Some(len) if **len > 0 => true,
+            _ => false
+        };
+
         let nickel_req = request::Request::from_internal(req,
                                                          &self.0.shared_data);
 
-        let nickel_res = response::Response::from_internal(res,
-                                                           &self.0.templates,
-                                                           &self.0.shared_data);
+        let mut nickel_res = response::Response::from_internal(res,
+                                                               &self.0.templates,
+                                                               &self.0.shared_data);
+
+        if request_has_body {
+            nickel_res.on_send(move |res| { res.set(Connection::close()); });
+        };
 
         self.0.middleware_stack.invoke(nickel_req, nickel_res);
     }
diff --git a/tests/examples/chaining.rs b/tests/examples/chaining.rs
index 859272ed5a..d18dae0d0d 100644
--- a/tests/examples/chaining.rs
+++ b/tests/examples/chaining.rs
@@ -17,6 +17,51 @@ where F: Fn(&mut Response),
     }
 }
 
+#[test]
+fn issue_326() {
+    use hyper::Client;
+    use hyper::header::Connection;
+
+    // The issue is that if we have a body for the request, and it isn't read
+    // by the server, it bleeds into the next request if keep-alive is active
+    run_example("chaining", |port| {
+        // The client is backed by a pool of keep-alive connections
+        let client = Client::new();
+        let url = format!("http://localhost:{}{}", port, "/post");
+
+        // scope the requests so the connection gets released back to the pool
+        {
+            // Assert that keep alive is still true after a zero length request body
+            // so that we know we're not dropping *all* requests with a body
+            let ref mut res = client.post(&url)
+                                    .send()
+                                    .unwrap();
+            assert_eq!(false, res.headers.has::<Connection>());
+            assert_eq!("post", read_body_to_string(res));
+        }
+
+        {
+            // Now send a request with a body, which won't be read by the server
+            let ref mut res = client.post(&url)
+                                    .body("0123456789")
+                                    .send()
+                                    .unwrap();
+            // we could assert that `Connection: close` was given, but we may
+            // want to be more flexible in future and allow buffers up to a
+            // certain length to be drained.
+            assert_eq!("post", read_body_to_string(res));
+        }
+
+        // Ensure the next request works as intended with no failure due to
+        // a corrupted stream.
+        let ref mut res = client.post(&url)
+                                .body("0123456789")
+                                .send()
+                                .unwrap();
+        assert_eq!("post", read_body_to_string(res));
+    })
+}
+
 mod expect_200 {
     use super::with_paths_and_method;
     use util::*;