Keep alive working

Author: rylev

src/body.rs

@@ -7,7 +7,7 @@ use std::pin::Pin;
 
 /// A streaming HTTP body.
 pub struct Body<R: AsyncRead> {
-    reader: Option<R>,
+    reader: R,
     length: Option<usize>,
 }
 
@@ -26,15 +26,15 @@ impl<R: AsyncRead> Body<R> {
     /// `Transfer-Encoding: Chunked`.
     pub fn new(reader: R) -> Self {
         Self {
-            reader: Some(reader),
+            reader,
             length: None,
         }
     }
 
     /// Create a new empty body.
-    pub fn empty() -> Self {
+    pub fn empty(reader: R) -> Self {
         Self {
-            reader: None,
+            reader,
             length: Some(0),
         }
     }
@@ -50,7 +50,7 @@ impl<R: AsyncRead> Body<R> {
 }
 
 impl<R: AsyncRead + Unpin> Body<R> {
-    pub fn into_reader(self) -> Option<R> {
+    pub fn into_reader(self) -> R {
         self.reader
     }
 }
@@ -61,7 +61,7 @@ impl Body<io::Cursor<Vec<u8>>> {
     pub fn from_string(string: String) -> Self {
         Self {
             length: Some(string.len()),
-            reader: Some(Cursor::new(string.into_bytes())),
+            reader: Cursor::new(string.into_bytes()),
         }
     }
 
@@ -70,7 +70,7 @@ impl Body<io::Cursor<Vec<u8>>> {
     pub fn from_bytes(bytes: Vec<u8>) -> Self {
         Self {
             length: Some(bytes.len()),
-            reader: Some(Cursor::new(bytes)),
+            reader: Cursor::new(bytes),
         }
     }
 }
@@ -81,9 +81,6 @@ impl<R: AsyncRead + Unpin> AsyncRead for Body<R> {
         cx: &mut Context<'_>,
         buf: &mut [u8],
     ) -> Poll<io::Result<usize>> {
-        match self.reader.as_mut() {
-            None => Poll::Ready(Ok(0)),
-            Some(reader) => Pin::new(&mut *reader).poll_read(cx, buf),
-        }
+        Pin::new(&mut self.reader).poll_read(cx, buf)
     }
 }

src/client.rs

@@ -129,7 +129,7 @@ where
         .find(|h| h.name == "Content-Length")
     {
         Some(_header) => Body::new(reader), // TODO: use the header value
-        None => Body::empty(),
+        None => Body::empty(reader),
     };
 
     // Return the response.

src/server.rs

@@ -13,18 +13,20 @@ use std::pin::Pin;
 
 use crate::{Body, Exception, MAX_HEADERS};
 
-pub async fn connect<'a, F, Fut, B>(
-    reader: &'a mut B,
-    writer: &'a mut B,
+pub async fn connect<'a, F, Fut, R, W, O: 'a>(
+    reader: &'a mut R,
+    writer: &'a mut W,
     callback: F,
 ) -> Result<(), Exception>
 where
-    F: Fn(&mut Request<Body<BufReader<&'a mut B>>>) -> Fut,
-    Fut: Future<Output = Result<Response<Body<&'a mut B>>, Exception>>,
-    B: Read + Write + Unpin + Send,
+    R: Read + Unpin + Send,
+    W: Write + Unpin,
+    F: Fn(&mut Request<Body<BufReader<&'a mut R>>>) -> Fut,
+    Fut: Future<Output = Result<Response<Body<O>>, Exception>>,
+    O: Read + Unpin + Send,
 {
     let req = decode(reader).await?;
-    if let OptionalRequest::Request(mut req) = req {
+    if let RequestOrReader::Request(mut req) = req {
         let headers = req.headers();
         let timeout = match (headers.get("Connection"), headers.get("Keep-Alive")) {
             (Some(connection), Some(_v))
@@ -38,22 +40,19 @@ where
 
         let beginning = Instant::now();
         loop {
-            println!("Handling request");
-            let mut res = encode(callback(&mut req).await?).await.unwrap();
+            // TODO: what to do when the callback returns Err
+            let mut res = encode(callback(&mut req).await?).await?;
             io::copy(&mut res, writer).await?;
-            let mut stream = res.body.into_reader().unwrap();
+            let mut stream = req.into_body().into_reader().into_inner();
             req = loop {
                 match decode(stream).await? {
-                    OptionalRequest::Request(r) => {
-                        break r;
-                    }
-                    OptionalRequest::Stream(r) => {
+                    RequestOrReader::Request(r) => break r,
+                    RequestOrReader::Reader(r) => {
                         let now = Instant::now();
                         if now - beginning > timeout {
                             return Ok(());
-                        } else {
-                            stream = r;
                         }
+                        stream = r;
                     }
                 }
             };
@@ -167,13 +166,13 @@ where
 }
 
 #[derive(Debug)]
-pub enum OptionalRequest<R: Read> {
+pub enum RequestOrReader<R: Read> {
     Request(Request<Body<BufReader<R>>>),
-    Stream(R),
+    Reader(R),
 }
 
 /// Decode an HTTP request on the server.
-pub async fn decode<R>(reader: R) -> Result<OptionalRequest<R>, Exception>
+pub async fn decode<R>(reader: R) -> Result<RequestOrReader<R>, Exception>
 where
     R: Read + Unpin + Send,
 {
@@ -187,7 +186,7 @@ where
         let bytes_read = reader.read_until(b'\n', &mut buf).await?;
         // No more bytes are yielded from the stream.
         if bytes_read == 0 {
-            return Ok(OptionalRequest::Stream(reader.into_inner()));
+            return Ok(RequestOrReader::Reader(reader.into_inner()));
         }
 
         // We've hit the end delimiter of the stream.
@@ -229,9 +228,9 @@ where
         .find(|h| h.name == "Content-Length")
     {
         Some(_header) => Body::new(reader), // TODO: use the header value
-        None => Body::empty(),
+        None => Body::empty(reader),
     };
 
     // Return the request.
-    Ok(OptionalRequest::Request(req.body(body)?))
+    Ok(RequestOrReader::Request(req.body(body)?))
 }

tests/test.rs

@@ -17,7 +17,7 @@ fn server() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
                 let (reader, writer) = &mut (&stream, &stream);
                 let req = server::decode(reader).await?;
                 dbg!(req);
-                let body: async_h1::Body<&[u8]> = async_h1::Body::empty();
+                let body: async_h1::Body<&[u8]> = async_h1::Body::empty(&[]);
                 let mut res = server::encode(http::Response::new(body)).await?;
                 io::copy(&mut res, writer).await?;
             }