From 042dcadf4323f001c83d4605da05454c9aa65018 Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Thu, 5 Jan 2017 15:13:13 -0600
Subject: [PATCH] Expose DataStream directly to allow for stream composition.

---
 lib/src/data/data.rs        | 18 ++++++++++--------
 lib/src/data/data_stream.rs | 17 +++++++++++++++--
 lib/src/data/mod.rs         | 26 ++++++++++++++++++++------
 lib/src/data/test_data.rs   | 26 +++++++++++++++++++++++---
 4 files changed, 68 insertions(+), 19 deletions(-)

diff --git a/lib/src/data/data.rs b/lib/src/data/data.rs
index 52021420..ad4790ef 100644
--- a/lib/src/data/data.rs
+++ b/lib/src/data/data.rs
@@ -1,4 +1,4 @@
-use std::io::{self, BufRead, Read, Write, Cursor, BufReader};
+use std::io::{self, Read, Write, Cursor, BufReader};
 use std::path::Path;
 use std::fs::File;
 use std::time::Duration;
@@ -61,7 +61,7 @@ impl Data {
     /// including that in the `peek` buffer. The method consumes the `Data`
     /// instance. This ensures that a `Data` type _always_ represents _all_ of
     /// the data in a request.
-    pub fn open(mut self) -> impl BufRead {
+    pub fn open(mut self) -> DataStream {
         // Swap out the buffer and stream for empty ones so we can move.
         let mut buffer = vec![];
         let mut stream = EmptyReader(self.stream.get_ref().clone());
@@ -71,13 +71,15 @@ impl Data {
         // Setup the underlying reader at the correct pointers.
         let mut cursor = Cursor::new(buffer);
         cursor.set_position(self.position as u64);
-        let buffered = cursor.take((self.capacity - self.position) as u64);
 
-        // Create the actual DataSteam.
-        DataStream {
-            network: stream.get_ref().clone(),
-            stream: buffered.chain(BufReader::new(stream)),
-        }
+        // Get a reference to the underlying network stream.
+        let network = stream.get_ref().clone();
+
+        // The first part of the stream is the buffer. Then the real steam.
+        let buf = cursor.take((self.capacity - self.position) as u64);
+        let stream = buf.chain(BufReader::new(stream));
+
+        DataStream::new(stream, network)
     }
 
     #[doc(hidden)]
diff --git a/lib/src/data/data_stream.rs b/lib/src/data/data_stream.rs
index b44a0399..1f7f983a 100644
--- a/lib/src/data/data_stream.rs
+++ b/lib/src/data/data_stream.rs
@@ -7,9 +7,22 @@ use http::hyper::h1::HttpReader;
 pub type StreamReader = HttpReader<HttpStream>;
 pub type InnerStream = Chain<Take<Cursor<Vec<u8>>>, BufReader<StreamReader>>;
 
+/// Raw data stream of a request body.
+///
+/// This stream can only be obtained by calling
+/// [Data::open](/rocket/data/struct.Data.html#method.open). The stream contains
+/// all of the data in the body of the request. It exposes no methods directly.
+/// Instead, it must be used as an opaque `Read` or `BufRead` structure.
 pub struct DataStream {
-    pub stream: InnerStream,
-    pub network: HttpStream,
+    stream: InnerStream,
+    network: HttpStream,
+}
+
+impl DataStream {
+    #[doc(hidden)]
+    pub fn new(stream: InnerStream, network: HttpStream) -> DataStream {
+        DataStream { stream: stream, network: network, }
+    }
 }
 
 impl Read for DataStream {
diff --git a/lib/src/data/mod.rs b/lib/src/data/mod.rs
index 6ed27509..af868a66 100644
--- a/lib/src/data/mod.rs
+++ b/lib/src/data/mod.rs
@@ -1,11 +1,25 @@
 //! Types and traits for reading and parsing request body data.
 
-#[cfg(any(test, feature = "testing"))] mod test_data;
-#[cfg(not(any(test, feature = "testing")))] mod data;
-#[cfg(not(any(test, feature = "testing")))] mod data_stream;
+#[cfg(any(test, feature = "testing"))]
+#[path = "."]
+mod items {
+    mod test_data;
+
+    pub use self::test_data::Data;
+    pub use self::test_data::DataStream;
+}
+
+#[cfg(not(any(test, feature = "testing")))]
+#[path = "."]
+mod items {
+    mod data;
+    mod data_stream;
+
+    pub use self::data::Data;
+    pub use self::data_stream::DataStream;
+}
+
 mod from_data;
 
 pub use self::from_data::{FromData, Outcome};
-
-#[cfg(any(test, feature = "testing"))] pub use self::test_data::Data;
-#[cfg(not(any(test, feature = "testing")))] pub use self::data::Data;
+pub use self::items::{Data, DataStream};
diff --git a/lib/src/data/test_data.rs b/lib/src/data/test_data.rs
index 6a0802e4..3ecb24b6 100644
--- a/lib/src/data/test_data.rs
+++ b/lib/src/data/test_data.rs
@@ -1,4 +1,4 @@
-use std::io::{self, BufRead, Write, Cursor, BufReader};
+use std::io::{self, Read, BufRead, Write, Cursor, BufReader};
 use std::path::Path;
 use std::fs::File;
 
@@ -11,13 +11,33 @@ pub type BodyReader<'a, 'b> =
 
 const PEEK_BYTES: usize = 4096;
 
+pub struct DataStream {
+    stream: BufReader<Cursor<Vec<u8>>>,
+}
+
+impl Read for DataStream {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.stream.read(buf)
+    }
+}
+
+impl BufRead for DataStream {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        self.stream.fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.stream.consume(amt)
+    }
+}
+
 pub struct Data {
     data: Vec<u8>,
 }
 
 impl Data {
-    pub fn open(self) -> impl BufRead {
-        BufReader::new(Cursor::new(self.data))
+    pub fn open(self) -> DataStream {
+        DataStream { stream: BufReader::new(Cursor::new(self.data)) }
     }
 
     #[inline(always)]