use crate::{Buffer, Conn, Error, HttpContext};
use std::{
pin::pin,
task::{Context, Poll, Waker},
};
use trillium_testing::{TestTransport, harness, test};
fn noop_context() -> Context<'static> {
Context::from_waker(Waker::noop())
}
#[test(harness)]
async fn dripped_bytes_do_not_multiply_the_head_buffer() {
const DRIPS: usize = 16;
let (client, mut server) = TestTransport::new();
let context = HttpContext::new();
let mut buffer = Buffer::with_capacity(32);
{
let mut head = pin!(Conn::head(&mut server, &mut buffer, &context));
let mut cx = noop_context();
for received in 0..DRIPS {
client.write_all(b"A");
assert!(
matches!(head.as_mut().poll(&mut cx), Poll::Pending),
"head completed after {received} dripped bytes, before any terminator"
);
}
}
let ceiling = 2 * (context.config.request_buffer_initial_len + DRIPS);
assert!(
buffer.capacity() <= ceiling,
"{DRIPS} dripped bytes grew the head buffer to {} bytes (ceiling {ceiling})",
buffer.capacity()
);
}
#[test(harness)]
async fn byte_at_a_time_head_still_parses() {
let request = b"GET /path?q HTTP/1.1\r\nHost: example\r\nAccept: */*\r\n\r\n";
let (client, mut server) = TestTransport::new();
let context = HttpContext::new();
let mut buffer = Buffer::from(Vec::with_capacity(16));
client.write_all(request);
let outcome = {
let mut head = pin!(Conn::head(&mut server, &mut buffer, &context));
let mut cx = noop_context();
let mut ready = None;
for _ in 0..100 {
if let Poll::Ready(result) = head.as_mut().poll(&mut cx) {
ready = Some(result);
break;
}
}
ready.expect("head never completed despite the full request being available")
};
let (head_size, _start_time) =
outcome.expect("a complete head dripped one byte at a time should parse");
assert_eq!(head_size, request.len());
assert_eq!(&buffer[..], &request[..]);
}
#[test(harness)]
async fn dribbled_incomplete_head_errors_with_bounded_capacity() {
let (client, mut server) = TestTransport::new();
let mut context = HttpContext::new();
context.config.head_max_len = 16;
let allowance = context.config.head_max_len;
let initial_capacity = 64;
let mut buffer = Buffer::with_capacity(initial_capacity);
let outcome = {
let mut head = pin!(Conn::head(&mut server, &mut buffer, &context));
let mut cx = noop_context();
let mut ready = None;
for _ in 0..allowance {
client.write_all(b"A");
if let Poll::Ready(result) = head.as_mut().poll(&mut cx) {
ready = Some(result);
break;
}
}
ready
};
assert!(
matches!(outcome, Some(Err(Error::HeadersTooLong))),
"expected HeadersTooLong after {allowance} dribbled bytes, got {outcome:?}"
);
let ceiling = 2 * (allowance + initial_capacity);
assert!(
buffer.capacity() <= ceiling,
"head buffer grew to {} bytes for a {allowance}-byte allowance",
buffer.capacity()
);
}