Merge pull request #10 from yoshuawuyts/keep-alive-server

yoshuawuyts · web-flow · commit 3fed77b33eb8 · 2019-11-12T19:12:36.000+01:00
Basic KeepAlive Server Support
diff --git a/Cargo.toml b/Cargo.toml
@@ -15,7 +15,7 @@ edition = "2018"
 httparse = "1.3.3"
 http = "0.1.17"
 futures-io = "0.3.0"
-async-std = "0.99.12"
+async-std = { version = "0.99.12", features = ["unstable"] }
 futures-core-preview = "0.3.0-alpha.18"
 
 [dev-dependencies]
diff --git a/examples/server.rs b/examples/server.rs
@@ -1,6 +1,7 @@
 use async_h1::{server, Body};
 use async_std::prelude::*;
-use async_std::{io, net, task};
+use async_std::net;
+use async_std::task;
 
 fn main() -> Result<(), async_h1::Exception> {
     task::block_on(async {
@@ -9,17 +10,14 @@ fn main() -> Result<(), async_h1::Exception> {
         let mut incoming = listener.incoming();
 
         while let Some(stream) = incoming.next().await {
-            task::spawn(async move {
+            task::spawn(async {
                 let stream = stream?;
                 let (reader, writer) = &mut (&stream, &stream);
-                let req = server::decode(reader).await?;
-                if req.is_some() {
-                    // dbg!(req);
+                server::connect(reader, writer, |_| async {
                     let body = Body::from_string("hello chashu".to_owned());
-                    let mut res = server::encode(http::Response::new(body)).await?;
-                    io::copy(&mut res, writer).await?;
-                }
-                Ok::<(), async_h1::Exception>(())
+                    Ok(http::Response::new(body))
+                })
+                .await
             });
         }
         Ok(())
diff --git a/src/body.rs b/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),
         }
     }
@@ -49,13 +49,19 @@ impl<R: AsyncRead> Body<R> {
     }
 }
 
+impl<R: AsyncRead + Unpin> Body<R> {
+    pub fn into_reader(self) -> R {
+        self.reader
+    }
+}
+
 impl Body<io::Cursor<Vec<u8>>> {
     /// Create a new instance from a string.
     #[inline]
     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()),
         }
     }
 
@@ -64,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),
         }
     }
 }
@@ -75,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)
     }
 }
diff --git a/src/client.rs b/src/client.rs
@@ -64,7 +64,6 @@ pub async fn encode<R: AsyncRead>(req: Request<Body<R>>) -> Result<Encoder<R>, s
         // See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Transfer-Encoding
         //      https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Trailer
     }
-
     for (header, value) in req.headers() {
         write!(
             &mut buf,
@@ -130,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.
diff --git a/src/server.rs b/src/server.rs
@@ -1,21 +1,72 @@
 //! Process HTTP connections on the server.
 
+use async_std::future::Future;
 use async_std::io::{self, BufReader};
+use async_std::io::{Read, Write};
 use async_std::prelude::*;
 use async_std::task::{Context, Poll};
 use futures_core::ready;
-use futures_io::AsyncRead;
 use http::{Request, Response, Version};
+use std::time::{Duration, Instant};
 
 use std::pin::Pin;
 
 use crate::{Body, Exception, MAX_HEADERS};
 
+pub async fn connect<'a, F, Fut, R, W, O: 'a>(
+    reader: &'a mut R,
+    writer: &'a mut W,
+    callback: F,
+) -> Result<(), Exception>
+where
+    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 RequestOrReader::Request(mut req) = req {
+        let headers = req.headers();
+        let timeout = match (headers.get("Connection"), headers.get("Keep-Alive")) {
+            (Some(connection), Some(_v))
+                if connection == http::header::HeaderValue::from_static("Keep-Alive") =>
+            {
+                // TODO: parse timeout
+                Duration::from_secs(5)
+            }
+            _ => Duration::from_secs(5),
+        };
+
+        let beginning = Instant::now();
+        loop {
+            // 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 = req.into_body().into_reader().into_inner();
+            req = loop {
+                match decode(stream).await? {
+                    RequestOrReader::Request(r) => break r,
+                    RequestOrReader::Reader(r) => {
+                        let now = Instant::now();
+                        if now - beginning > timeout {
+                            return Ok(());
+                        }
+                        stream = r;
+                    }
+                }
+            };
+        }
+    }
+
+    Ok(())
+}
+
 /// A streaming HTTP encoder.
 ///
 /// This is returned from [`encode`].
 #[derive(Debug)]
-pub struct Encoder<R: AsyncRead> {
+pub struct Encoder<R: Read> {
     /// Keep track how far we've indexed into the headers + body.
     cursor: usize,
     /// HTTP headers to be sent.
@@ -30,7 +81,7 @@ pub struct Encoder<R: AsyncRead> {
     body_bytes_read: usize,
 }
 
-impl<R: AsyncRead> Encoder<R> {
+impl<R: Read> Encoder<R> {
     /// Create a new instance.
     pub(crate) fn new(headers: Vec<u8>, body: Body<R>) -> Self {
         Self {
@@ -44,7 +95,7 @@ impl<R: AsyncRead> Encoder<R> {
     }
 }
 
-impl<R: AsyncRead + Unpin> AsyncRead for Encoder<R> {
+impl<R: Read + Unpin> Read for Encoder<R> {
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -81,43 +132,49 @@ impl<R: AsyncRead + Unpin> AsyncRead for Encoder<R> {
 // TODO: return a reader in the response
 pub async fn encode<R>(res: Response<Body<R>>) -> io::Result<Encoder<R>>
 where
-    R: AsyncRead,
+    R: Read + Send,
 {
     let mut buf: Vec<u8> = vec![];
 
     let reason = res.status().canonical_reason().unwrap();
     let status = res.status();
-    write!(&mut buf, "HTTP/1.1 {} {}\r\n", status.as_str(), reason).await?;
+    std::io::Write::write_fmt(
+        &mut buf,
+        format_args!("HTTP/1.1 {} {}\r\n", status.as_str(), reason),
+    )?;
 
     // If the body isn't streaming, we can set the content-length ahead of time. Else we need to
     // send all items in chunks.
     if let Some(len) = res.body().len() {
-        write!(&mut buf, "Content-Length: {}\r\n", len).await?;
+        std::io::Write::write_fmt(&mut buf, format_args!("Content-Length: {}\r\n", len))?;
     } else {
-        write!(&mut buf, "Transfer-Encoding: chunked\r\n").await?;
+        std::io::Write::write_fmt(&mut buf, format_args!("Transfer-Encoding: chunked\r\n"))?;
         panic!("chunked encoding is not implemented yet");
         // See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Transfer-Encoding
         //      https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Trailer
     }
 
     for (header, value) in res.headers() {
-        write!(
+        std::io::Write::write_fmt(
             &mut buf,
-            "{}: {}\r\n",
-            header.as_str(),
-            value.to_str().unwrap()
-        )
-        .await?;
+            format_args!("{}: {}\r\n", header.as_str(), value.to_str().unwrap()),
+        )?
     }
 
-    write!(&mut buf, "\r\n").await?;
+    std::io::Write::write_fmt(&mut buf, format_args!("\r\n"))?;
     Ok(Encoder::new(buf, res.into_body()))
 }
 
+#[derive(Debug)]
+pub enum RequestOrReader<R: Read> {
+    Request(Request<Body<BufReader<R>>>),
+    Reader(R),
+}
+
 /// Decode an HTTP request on the server.
-pub async fn decode<R>(reader: R) -> Result<Option<Request<Body<BufReader<R>>>>, Exception>
+pub async fn decode<R>(reader: R) -> Result<RequestOrReader<R>, Exception>
 where
-    R: AsyncRead + Unpin + Send,
+    R: Read + Unpin + Send,
 {
     let mut reader = BufReader::new(reader);
     let mut buf = Vec::new();
@@ -129,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(None);
+            return Ok(RequestOrReader::Reader(reader.into_inner()));
         }
 
         // We've hit the end delimiter of the stream.
@@ -171,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(Some(req.body(body)?))
+    Ok(RequestOrReader::Request(req.body(body)?))
 }
diff --git a/tests/test.rs b/tests/test.rs
@@ -17,9 +17,8 @@ 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 res = http::Response::new(body);
-                let mut res = server::encode(res).await?;
+                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?;
             }
 

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ use std::pin::Pin;`
`7`	`7`
`8`	`8`	`/// A streaming HTTP body.`
`9`	`9`	`pub struct Body<R: AsyncRead> {`
`10`		`- reader: Option<R>,`
	`10`	`+ reader: R,`
`11`	`11`	`length: Option<usize>,`
`12`	`12`	`}`
`13`	`13`
`@@ -26,15 +26,15 @@ impl<R: AsyncRead> Body<R> {`
`26`	`26`	/// `Transfer-Encoding: Chunked`.
`27`	`27`	`pub fn new(reader: R) -> Self {`
`28`	`28`	`Self {`
`29`		`- reader: Some(reader),`
	`29`	`+ reader,`
`30`	`30`	`length: None,`
`31`	`31`	`}`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`/// Create a new empty body.`
`35`		`- pub fn empty() -> Self {`
	`35`	`+ pub fn empty(reader: R) -> Self {`
`36`	`36`	`Self {`
`37`		`- reader: None,`
	`37`	`+ reader,`
`38`	`38`	`length: Some(0),`
`39`	`39`	`}`
`40`	`40`	`}`
`@@ -49,13 +49,19 @@ impl<R: AsyncRead> Body<R> {`
`49`	`49`	`}`
`50`	`50`	`}`
`51`	`51`
	`52`	`+impl<R: AsyncRead + Unpin> Body<R> {`
	`53`	`+ pub fn into_reader(self) -> R {`
	`54`	`+ self.reader`
	`55`	`+ }`
	`56`	`+}`
	`57`	`+`
`52`	`58`	`impl Body<io::Cursor<Vec<u8>>> {`
`53`	`59`	`/// Create a new instance from a string.`
`54`	`60`	`#[inline]`
`55`	`61`	`pub fn from_string(string: String) -> Self {`
`56`	`62`	`Self {`
`57`	`63`	`length: Some(string.len()),`
`58`		`- reader: Some(Cursor::new(string.into_bytes())),`
	`64`	`+ reader: Cursor::new(string.into_bytes()),`
`59`	`65`	`}`
`60`	`66`	`}`
`61`	`67`
`@@ -64,7 +70,7 @@ impl Body<io::Cursor<Vec<u8>>> {`
`64`	`70`	`pub fn from_bytes(bytes: Vec<u8>) -> Self {`
`65`	`71`	`Self {`
`66`	`72`	`length: Some(bytes.len()),`
`67`		`- reader: Some(Cursor::new(bytes)),`
	`73`	`+ reader: Cursor::new(bytes),`
`68`	`74`	`}`
`69`	`75`	`}`
`70`	`76`	`}`
`@@ -75,9 +81,6 @@ impl<R: AsyncRead + Unpin> AsyncRead for Body<R> {`
`75`	`81`	`cx: &mut Context<'_>,`
`76`	`82`	`buf: &mut [u8],`
`77`	`83`	`) -> Poll<io::Result<usize>> {`
`78`		`- match self.reader.as_mut() {`
`79`		`- None => Poll::Ready(Ok(0)),`
`80`		`- Some(reader) => Pin::new(&mut *reader).poll_read(cx, buf),`
`81`		`- }`
	`84`	`+ Pin::new(&mut self.reader).poll_read(cx, buf)`
`82`	`85`	`}`
`83`	`86`	`}`