keep-alive: always read body to end

Author: jbr

Cargo.toml

@@ -27,12 +27,8 @@ futures-core = "0.3.8"
 log = "0.4.11"
 pin-project = "1.0.2"
 async-channel = "1.5.1"
+async-dup = "1.2.2"
 
 [dev-dependencies]
 pretty_assertions = "0.6.1"
 async-std = { version = "1.7.0", features = ["attributes"] }
-tempfile = "3.1.0"
-async-test = "1.0.0"
-duplexify = "1.2.2"
-async-dup = "1.2.2"
-async-channel = "1.5.1"

src/chunked/decoder.rs

@@ -19,7 +19,8 @@ lazy_static::lazy_static! {
 
 /// Decodes a chunked body according to
 /// https://tools.ietf.org/html/rfc7230#section-4.1
-pub(crate) struct ChunkedDecoder<R: Read> {
+#[derive(Debug)]
+pub struct ChunkedDecoder<R: Read> {
     /// The underlying stream
     inner: R,
     /// Buffer for the already read, but not yet parsed data.

src/read_notifier.rs

@@ -26,7 +26,7 @@ impl<B> fmt::Debug for ReadNotifier<B> {
     }
 }
 
-impl<B: BufRead> ReadNotifier<B> {
+impl<B: Read> ReadNotifier<B> {
     pub(crate) fn new(reader: B, sender: Sender<()>) -> Self {
         Self {
             reader,

src/server/body_reader.rs

@@ -0,0 +1,35 @@
+use crate::chunked::ChunkedDecoder;
+use async_dup::{Arc, Mutex};
+use async_std::io::{BufReader, Read, Take};
+use async_std::task::{Context, Poll};
+use std::{fmt::Debug, io, pin::Pin};
+
+pub enum BodyReader<IO: Read + Unpin> {
+    Chunked(Arc<Mutex<ChunkedDecoder<BufReader<IO>>>>),
+    Fixed(Arc<Mutex<Take<BufReader<IO>>>>),
+    None,
+}
+
+impl<IO: Read + Unpin> Debug for BodyReader<IO> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            BodyReader::Chunked(_) => f.write_str("BodyReader::Chunked"),
+            BodyReader::Fixed(_) => f.write_str("BodyReader::Fixed"),
+            BodyReader::None => f.write_str("BodyReader::None"),
+        }
+    }
+}
+
+impl<IO: Read + Unpin> Read for BodyReader<IO> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        match &*self {
+            BodyReader::Chunked(r) => Pin::new(&mut *r.lock()).poll_read(cx, buf),
+            BodyReader::Fixed(r) => Pin::new(&mut *r.lock()).poll_read(cx, buf),
+            BodyReader::None => Poll::Ready(Ok(0)),
+        }
+    }
+}

src/server/decode.rs

@@ -2,12 +2,14 @@
 
 use std::str::FromStr;
 
+use async_dup::{Arc, Mutex};
 use async_std::io::{BufReader, Read, Write};
 use async_std::{prelude::*, task};
 use http_types::headers::{CONTENT_LENGTH, EXPECT, TRANSFER_ENCODING};
 use http_types::{ensure, ensure_eq, format_err};
 use http_types::{Body, Method, Request, Url};
 
+use super::body_reader::BodyReader;
 use crate::chunked::ChunkedDecoder;
 use crate::read_notifier::ReadNotifier;
 use crate::{MAX_HEADERS, MAX_HEAD_LENGTH};
@@ -21,7 +23,7 @@ const CONTINUE_HEADER_VALUE: &str = "100-continue";
 const CONTINUE_RESPONSE: &[u8] = b"HTTP/1.1 100 Continue\r\n\r\n";
 
 /// Decode an HTTP request on the server.
-pub async fn decode<IO>(mut io: IO) -> http_types::Result<Option<Request>>
+pub async fn decode<IO>(mut io: IO) -> http_types::Result<Option<(Request, BodyReader<IO>)>>
 where
     IO: Read + Write + Clone + Send + Sync + Unpin + 'static,
 {
@@ -108,26 +110,29 @@ where
     }
 
     // Check for Transfer-Encoding
-    if let Some(encoding) = transfer_encoding {
-        if encoding.last().as_str() == "chunked" {
-            let trailer_sender = req.send_trailers();
-            let reader = ChunkedDecoder::new(reader, trailer_sender);
-            let reader = BufReader::new(reader);
-            let reader = ReadNotifier::new(reader, body_read_sender);
-            req.set_body(Body::from_reader(reader, None));
-            return Ok(Some(req));
-        }
-        // Fall through to Content-Length
-    }
-
-    // Check for Content-Length.
-    if let Some(len) = content_length {
+    if transfer_encoding
+        .map(|te| te.as_str().eq_ignore_ascii_case("chunked"))
+        .unwrap_or(false)
+    {
+        let trailer_sender = req.send_trailers();
+        let reader = ChunkedDecoder::new(reader, trailer_sender);
+        let reader = Arc::new(Mutex::new(reader));
+        let reader_clone = reader.clone();
+        let reader = ReadNotifier::new(reader, body_read_sender);
+        let reader = BufReader::new(reader);
+        req.set_body(Body::from_reader(reader, None));
+        return Ok(Some((req, BodyReader::Chunked(reader_clone))));
+    } else if let Some(len) = content_length {
         let len = len.last().as_str().parse::<usize>()?;
-        let reader = ReadNotifier::new(reader.take(len as u64), body_read_sender);
-        req.set_body(Body::from_reader(reader, Some(len)));
+        let reader = Arc::new(Mutex::new(reader.take(len as u64)));
+        req.set_body(Body::from_reader(
+            BufReader::new(ReadNotifier::new(reader.clone(), body_read_sender)),
+            Some(len),
+        ));
+        Ok(Some((req, BodyReader::Fixed(reader))))
+    } else {
+        Ok(Some((req, BodyReader::None)))
     }
-
-    Ok(Some(req))
 }
 
 fn url_from_httparse_req(req: &httparse::Request<'_, '_>) -> http_types::Result<Url> {

src/server/mod.rs

@@ -1,13 +1,12 @@
 //! Process HTTP connections on the server.
 
-use std::time::Duration;
-
 use async_std::future::{timeout, Future, TimeoutError};
 use async_std::io::{self, Read, Write};
 use http_types::headers::{CONNECTION, UPGRADE};
 use http_types::upgrade::Connection;
 use http_types::{Request, Response, StatusCode};
-
+use std::{marker::PhantomData, time::Duration};
+mod body_reader;
 mod decode;
 mod encode;
 
@@ -38,14 +37,14 @@ where
     F: Fn(Request) -> Fut,
     Fut: Future<Output = http_types::Result<Response>>,
 {
-    accept_with_opts(io, endpoint, Default::default()).await
+    Server::new(io, endpoint).accept().await
 }
 
 /// Accept a new incoming HTTP/1.1 connection.
 ///
 /// Supports `KeepAlive` requests by default.
 pub async fn accept_with_opts<RW, F, Fut>(
-    mut io: RW,
+    io: RW,
     endpoint: F,
     opts: ServerOptions,
 ) -> http_types::Result<()>
@@ -54,35 +53,99 @@ where
     F: Fn(Request) -> Fut,
     Fut: Future<Output = http_types::Result<Response>>,
 {
-    loop {
+    Server::new(io, endpoint).with_opts(opts).accept().await
+}
+
+/// struct for server
+#[derive(Debug)]
+pub struct Server<RW, F, Fut> {
+    io: RW,
+    endpoint: F,
+    opts: ServerOptions,
+    _phantom: PhantomData<Fut>,
+}
+
+/// An enum that represents whether the server should accept a subsequent request
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum ConnectionStatus {
+    /// The server should not accept another request
+    Close,
+
+    /// The server may accept another request
+    KeepAlive,
+}
+
+impl<RW, F, Fut> Server<RW, F, Fut>
+where
+    RW: Read + Write + Clone + Send + Sync + Unpin + 'static,
+    F: Fn(Request) -> Fut,
+    Fut: Future<Output = http_types::Result<Response>>,
+{
+    /// builds a new server
+    pub fn new(io: RW, endpoint: F) -> Self {
+        Self {
+            io,
+            endpoint,
+            opts: Default::default(),
+            _phantom: PhantomData,
+        }
+    }
+
+    /// with opts
+    pub fn with_opts(mut self, opts: ServerOptions) -> Self {
+        self.opts = opts;
+        self
+    }
+
+    /// accept in a loop
+    pub async fn accept(&mut self) -> http_types::Result<()> {
+        while ConnectionStatus::KeepAlive == self.accept_one().await? {}
+        Ok(())
+    }
+
+    /// accept one request
+    pub async fn accept_one(&mut self) -> http_types::Result<ConnectionStatus>
+    where
+        RW: Read + Write + Clone + Send + Sync + Unpin + 'static,
+        F: Fn(Request) -> Fut,
+        Fut: Future<Output = http_types::Result<Response>>,
+    {
         // Decode a new request, timing out if this takes longer than the timeout duration.
-        let fut = decode(io.clone());
+        let fut = decode(self.io.clone());
 
-        let req = if let Some(timeout_duration) = opts.headers_timeout {
+        let (req, mut body) = if let Some(timeout_duration) = self.opts.headers_timeout {
             match timeout(timeout_duration, fut).await {
                 Ok(Ok(Some(r))) => r,
-                Ok(Ok(None)) | Err(TimeoutError { .. }) => break, /* EOF or timeout */
+                Ok(Ok(None)) | Err(TimeoutError { .. }) => return Ok(ConnectionStatus::Close), /* EOF or timeout */
                 Ok(Err(e)) => return Err(e),
             }
         } else {
             match fut.await? {
                 Some(r) => r,
-                None => break, /* EOF */
+                None => return Ok(ConnectionStatus::Close), /* EOF */
             }
         };
 
         let has_upgrade_header = req.header(UPGRADE).is_some();
-        let connection_header_is_upgrade = req
+        let connection_header_as_str = req
             .header(CONNECTION)
-            .map(|connection| connection.as_str().eq_ignore_ascii_case("upgrade"))
-            .unwrap_or(false);
+            .map(|connection| connection.as_str())
+            .unwrap_or("");
+
+        let connection_header_is_upgrade = connection_header_as_str.eq_ignore_ascii_case("upgrade");
+        let mut close_connection = connection_header_as_str.eq_ignore_ascii_case("close");
 
         let upgrade_requested = has_upgrade_header && connection_header_is_upgrade;
 
         let method = req.method();
 
         // Pass the request to the endpoint and encode the response.
-        let mut res = endpoint(req).await?;
+        let mut res = (self.endpoint)(req).await?;
+
+        close_connection |= res
+            .header(CONNECTION)
+            .map(|c| c.as_str().eq_ignore_ascii_case("close"))
+            .unwrap_or(false);
 
         let upgrade_provided = res.status() == StatusCode::SwitchingProtocols && res.has_upgrade();
 
@@ -94,14 +157,22 @@ where
 
         let mut encoder = Encoder::new(res, method);
 
-        // Stream the response to the writer.
-        io::copy(&mut encoder, &mut io).await?;
+        let bytes_written = io::copy(&mut encoder, &mut self.io).await?;
+        log::trace!("wrote {} response bytes", bytes_written);
+
+        let body_bytes_discarded = io::copy(&mut body, &mut io::sink()).await?;
+        log::trace!(
+            "discarded {} unread request body bytes",
+            body_bytes_discarded
+        );
 
         if let Some(upgrade_sender) = upgrade_sender {
-            upgrade_sender.send(Connection::new(io.clone())).await;
-            return Ok(());
+            upgrade_sender.send(Connection::new(self.io.clone())).await;
+            return Ok(ConnectionStatus::Close);
+        } else if close_connection {
+            Ok(ConnectionStatus::Close)
+        } else {
+            Ok(ConnectionStatus::KeepAlive)
         }
     }
-
-    Ok(())
 }