Skip to content

Commit

Permalink
don't panic if the connection is closed at any point
Browse files Browse the repository at this point in the history
  • Loading branch information
jbr committed Mar 3, 2021
1 parent 9993b04 commit dffae5f
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 8 deletions.
7 changes: 6 additions & 1 deletion src/client/decode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,12 @@ where
loop {
let bytes_read = reader.read_until(LF, &mut buf).await?;
// No more bytes are yielded from the stream.
assert!(bytes_read != 0, "Empty response"); // TODO: ensure?

match (bytes_read, buf.len()) {
(0, 0) => return Err(format_err!("connection closed")),
(0, _) => return Err(format_err!("empty response")),
_ => {}
}

// Prevent CWE-400 DDOS with large HTTP Headers.
ensure!(
Expand Down
25 changes: 25 additions & 0 deletions tests/client_decode.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,9 @@
mod test_utils;

mod client_decode {
use std::io::Write;

use super::test_utils::CloseableCursor;
use async_h1::client;
use async_std::io::Cursor;
use http_types::headers;
Expand Down Expand Up @@ -42,6 +47,26 @@ mod client_decode {
Ok(())
}

#[async_std::test]
async fn connection_closure() -> Result<()> {
let mut cursor = CloseableCursor::default();
cursor.write_all(b"HTTP/1.1 200 OK\r\nhost: example.com")?;
cursor.close();
assert_eq!(
client::decode(cursor).await.unwrap_err().to_string(),
"empty response"
);

let cursor = CloseableCursor::default();
cursor.close();
assert_eq!(
client::decode(cursor).await.unwrap_err().to_string(),
"connection closed"
);

Ok(())
}

#[async_std::test]
async fn response_newlines() -> Result<()> {
let res = decode_lines(vec![
Expand Down
34 changes: 27 additions & 7 deletions tests/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use async_h1::{
client::Encoder,
server::{ConnectionStatus, Server},
};
use async_std::io::{Read, Write};
use async_std::io::{Read as AsyncRead, Write as AsyncWrite};
use http_types::{Request, Response, Result};
use std::{
fmt::{Debug, Display},
Expand Down Expand Up @@ -58,7 +58,7 @@ where
}
}

impl<F, Fut> Read for TestServer<F, Fut>
impl<F, Fut> AsyncRead for TestServer<F, Fut>
where
F: Fn(Request) -> Fut,
Fut: Future<Output = Result<Response>>,
Expand All @@ -72,7 +72,7 @@ where
}
}

impl<F, Fut> Write for TestServer<F, Fut>
impl<F, Fut> AsyncWrite for TestServer<F, Fut>
where
F: Fn(Request) -> Fut,
Fut: Future<Output = Result<Response>>,
Expand Down Expand Up @@ -187,7 +187,17 @@ impl Debug for CloseableCursor {
}
}

impl Read for &CloseableCursor {
impl AsyncRead for CloseableCursor {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut &*self).poll_read(cx, buf)
}
}

impl AsyncRead for &CloseableCursor {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
Expand All @@ -209,7 +219,7 @@ impl Read for &CloseableCursor {
}
}

impl Write for &CloseableCursor {
impl AsyncWrite for &CloseableCursor {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
Expand Down Expand Up @@ -237,7 +247,7 @@ impl Write for &CloseableCursor {
}
}

impl Read for TestIO {
impl AsyncRead for TestIO {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
Expand All @@ -247,7 +257,7 @@ impl Read for TestIO {
}
}

impl Write for TestIO {
impl AsyncWrite for TestIO {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
Expand All @@ -264,3 +274,13 @@ impl Write for TestIO {
Pin::new(&mut &*self.write).poll_close(cx)
}
}

impl std::io::Write for CloseableCursor {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.write().unwrap().data.write(buf)
}

fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}

0 comments on commit dffae5f

Please sign in to comment.