use alloc::vec::Vec;
use crate::{
body::{Decoder, Encoder, Step, push_framing},
buffer::Buffer,
error::HttpError,
head::{
Asked, HeadKind, HeadScanner, Progress, ResponseHead, ResponseInfo, interpret_response,
},
limits::Limits,
message::{OutBody, Request},
method::Method,
};
#[non_exhaustive]
#[derive(Debug)]
pub enum ClientEvent<'a> {
Response(ResponseHead<'a>),
Body(&'a [u8]),
ResponseEnd,
NeedInput,
Idle,
Upgraded,
Closed,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
Idle,
Head,
Body,
Complete,
Upgraded,
Closed,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ReqWrite {
Streaming(Encoder),
Sent,
}
pub struct ClientConn {
buf: Buffer,
limits: Limits,
scanner: HeadScanner,
trailers: HeadScanner,
decoder: Decoder,
info: ResponseInfo,
asked: Asked,
upgrade_asked: bool,
state: State,
req: ReqWrite,
keep_alive: bool,
interim_bytes: usize,
}
impl Default for ClientConn {
fn default() -> Self {
Self::new()
}
}
impl core::fmt::Debug for ClientConn {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ClientConn")
.field("state", &self.state)
.field("request", &self.req)
.field("buffered", &self.buf.unread().len())
.field("limits", &self.limits)
.finish()
}
}
impl ClientConn {
#[must_use]
pub fn new() -> Self {
Self::with_limits(Limits::new())
}
#[must_use]
pub fn with_limits(limits: Limits) -> Self {
Self {
buf: Buffer::new(limits.buffer_capacity()),
limits,
scanner: HeadScanner::new(),
trailers: HeadScanner::new(),
decoder: Decoder::new(),
info: ResponseInfo::EMPTY,
asked: Asked {
head: false,
connect: false,
},
upgrade_asked: false,
state: State::Idle,
req: ReqWrite::Sent,
keep_alive: true,
interim_bytes: 0,
}
}
#[must_use]
pub fn limits(&self) -> &Limits {
&self.limits
}
#[must_use = "bytes beyond the returned count were not taken"]
pub fn feed(&mut self, bytes: &[u8]) -> usize {
self.buf.feed(bytes)
}
pub fn read_buf(&mut self) -> &mut [u8] {
self.buf.spare()
}
pub fn commit(&mut self, n: usize) {
self.buf.commit(n);
}
pub fn feed_eof(&mut self) {
self.buf.set_eof();
}
pub fn send_request(&mut self, request: &Request, out: &mut Vec<u8>) -> Result<(), HttpError> {
if let Some(e) = request.error() {
return Err(e);
}
match self.state {
State::Idle => {}
State::Closed | State::Upgraded => {
return Err(HttpError::InvalidState(
"the connection is closed or upgraded",
));
}
_ => return Err(HttpError::InvalidState("a request is already outstanding")),
}
if !self.buf.unread().is_empty() {
return Err(HttpError::UnexpectedData);
}
if !request.has_host {
return Err(HttpError::MissingHost);
}
let method = request.method_kind;
let expects_body = matches!(method, Method::Post | Method::Put | Method::Patch);
let (encoder, declared) = match &request.body {
OutBody::Full(b) if b.is_empty() && !expects_body => (Encoder::Discard, None),
OutBody::Full(b) => (
Encoder::Length {
remaining: b.len() as u64,
},
Some(b.len() as u64),
),
OutBody::Length(n) => (Encoder::Length { remaining: *n }, Some(*n)),
OutBody::Chunked => (Encoder::Chunked, None),
};
out.extend_from_slice(request.method().as_bytes());
out.push(b' ');
out.extend_from_slice(request.target().as_bytes());
out.extend_from_slice(b" HTTP/1.1\r\n");
out.extend_from_slice(request.fields());
push_framing(out, encoder, declared);
if request.close {
out.extend_from_slice(b"connection: close\r\n");
}
out.extend_from_slice(b"\r\n");
let mut encoder = encoder;
match &request.body {
OutBody::Full(b) => {
encoder.write(b, out)?;
self.req = ReqWrite::Sent;
}
_ => self.req = ReqWrite::Streaming(encoder),
}
self.asked = Asked {
head: method == Method::Head,
connect: method == Method::Connect,
};
self.upgrade_asked = request.upgrade;
self.keep_alive = !request.close;
self.buf.unpin();
self.buf.pin_at_pos();
self.scanner.reset();
self.interim_bytes = 0;
self.state = State::Head;
Ok(())
}
pub fn send_body(&mut self, data: &[u8], out: &mut Vec<u8>) -> Result<(), HttpError> {
match &mut self.req {
ReqWrite::Streaming(encoder) => encoder.write(data, out),
ReqWrite::Sent => Err(HttpError::InvalidState("no streamed request in progress")),
}
}
pub fn end_body(&mut self, out: &mut Vec<u8>) -> Result<(), HttpError> {
match &mut self.req {
ReqWrite::Streaming(encoder) => {
encoder.finish(out)?;
self.req = ReqWrite::Sent;
Ok(())
}
ReqWrite::Sent => Err(HttpError::InvalidState("no streamed request in progress")),
}
}
pub fn poll(&mut self) -> Result<ClientEvent<'_>, HttpError> {
loop {
match self.state {
State::Idle => {
if !self.buf.unread().is_empty() {
return Err(self.fail(HttpError::UnexpectedData));
}
if self.buf.is_eof() {
self.state = State::Closed;
return Ok(ClientEvent::Closed);
}
return Ok(ClientEvent::Idle);
}
State::Head => return self.poll_head(),
State::Body => return self.poll_body(),
State::Complete => {
let reusable =
self.keep_alive && self.info.keep_alive && self.req == ReqWrite::Sent;
if reusable {
self.state = State::Idle;
} else {
self.state = State::Closed;
return Ok(ClientEvent::Closed);
}
}
State::Upgraded => return Ok(ClientEvent::Upgraded),
State::Closed => return Ok(ClientEvent::Closed),
}
}
}
fn poll_head(&mut self) -> Result<ClientEvent<'_>, HttpError> {
loop {
let start = self.buf.pin();
let progress = self
.scanner
.advance(self.buf.from(start), &self.limits, HeadKind::Response, true)
.map_err(|e| self.fail(e))?;
let len = match progress {
Progress::Partial if self.buf.is_eof() => {
return Err(self.fail(HttpError::UnexpectedEof));
}
Progress::Partial => return Ok(ClientEvent::NeedInput),
Progress::Done(len) => len,
};
let info = interpret_response(
self.buf.from(start),
self.scanner.start_line(),
&self.scanner.fields,
&self.limits,
self.asked,
)
.map_err(|e| self.fail(e))?;
self.buf.consume(len);
self.buf.extend_pin_to_pos();
self.info = info;
let status = info.status;
let switched = if status == 101 {
if !self.upgrade_asked || !info.upgrade_field {
return Err(self.fail(HttpError::UnexpectedUpgrade));
}
true
} else {
self.asked.connect && (200..300).contains(&status)
};
if switched {
self.state = State::Upgraded;
} else if (100..200).contains(&status) {
self.interim_bytes = self.interim_bytes.saturating_add(len);
if self.interim_bytes > self.limits.max_header_bytes() {
return Err(self.fail(HttpError::HeadersTooLarge));
}
self.buf.unpin();
self.buf.pin_at_pos();
self.scanner.reset();
continue;
} else {
self.decoder.start(info.framing);
self.state = State::Body;
}
return Ok(ClientEvent::Response(self.view()));
}
}
fn poll_body(&mut self) -> Result<ClientEvent<'_>, HttpError> {
let mut framing_bytes = 0;
let step = self.decoder.step(
self.buf.unread(),
self.buf.is_eof(),
&self.limits,
&mut self.trailers,
&mut framing_bytes,
);
self.buf.consume(framing_bytes);
match step.map_err(|e| self.fail(e))? {
Step::Data(n) => {
let start = self.buf.pos();
self.buf.consume(n);
Ok(ClientEvent::Body(self.buf.range(start, start + n)))
}
Step::Done => {
self.state = State::Complete;
Ok(ClientEvent::ResponseEnd)
}
Step::NeedMore => Ok(ClientEvent::NeedInput),
}
}
fn fail(&mut self, e: HttpError) -> HttpError {
self.state = State::Closed;
e
}
fn view(&self) -> ResponseHead<'_> {
ResponseHead::new(
self.buf.from(self.buf.pin()),
&self.scanner.fields,
&self.info,
)
}
#[must_use]
pub fn response(&self) -> Option<ResponseHead<'_>> {
match self.state {
State::Body | State::Complete | State::Upgraded => Some(self.view()),
_ => None,
}
}
#[must_use]
pub fn is_upgraded(&self) -> bool {
self.state == State::Upgraded
}
#[must_use]
pub fn remaining_input(&self) -> &[u8] {
self.buf.unread()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::head::BodyFraming;
fn get() -> Request {
Request::get("/").header("host", "a")
}
#[test]
fn test_interim_responses_are_skipped() {
let mut conn = ClientConn::new();
conn.send_request(&get(), &mut Vec::new()).unwrap();
let _ = conn.feed(b"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 103 Early Hints\r\nLink: </a>\r\n\r\nHTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok");
let ClientEvent::Response(r) = conn.poll().unwrap() else {
panic!("expected response");
};
assert_eq!(r.status(), 200);
assert!(matches!(conn.poll().unwrap(), ClientEvent::Body(b"ok")));
assert!(matches!(conn.poll().unwrap(), ClientEvent::ResponseEnd));
assert!(matches!(conn.poll().unwrap(), ClientEvent::Idle));
}
#[test]
fn test_unrequested_upgrade_and_unsolicited_data_are_errors() {
let mut conn = ClientConn::new();
conn.send_request(&get(), &mut Vec::new()).unwrap();
let _ = conn
.feed(b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: x\r\nConnection: upgrade\r\n\r\n");
assert_eq!(conn.poll().unwrap_err(), HttpError::UnexpectedUpgrade);
let mut conn = ClientConn::new();
let _ = conn.feed(b"HTTP/1.1 200 OK\r\n\r\n");
assert_eq!(conn.poll().unwrap_err(), HttpError::UnexpectedData);
}
#[test]
fn test_head_response_has_no_body_and_connection_is_reused() {
let mut conn = ClientConn::new();
conn.send_request(
&Request::new("HEAD", "/").header("host", "a"),
&mut Vec::new(),
)
.unwrap();
let _ = conn.feed(b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n");
let ClientEvent::Response(r) = conn.poll().unwrap() else {
panic!("expected response");
};
assert_eq!(r.framing(), BodyFraming::None);
assert!(matches!(conn.poll().unwrap(), ClientEvent::ResponseEnd));
assert!(matches!(conn.poll().unwrap(), ClientEvent::Idle));
conn.send_request(&get(), &mut Vec::new()).unwrap();
assert!(matches!(
conn.send_request(&get(), &mut Vec::new()),
Err(HttpError::InvalidState(_))
));
}
#[test]
fn test_close_and_streaming_request_prevent_reuse() {
let mut conn = ClientConn::new();
conn.send_request(&get().close(), &mut Vec::new()).unwrap();
let _ = conn.feed(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n");
while !matches!(conn.poll().unwrap(), ClientEvent::ResponseEnd) {}
assert!(matches!(conn.poll().unwrap(), ClientEvent::Closed));
assert!(conn.send_request(&get(), &mut Vec::new()).is_err());
let mut conn = ClientConn::new();
let mut out = Vec::new();
conn.send_request(&Request::post("/").header("host", "a").chunked(), &mut out)
.unwrap();
let _ = conn.feed(b"HTTP/1.1 413 Content Too Large\r\nContent-Length: 0\r\n\r\n");
while !matches!(conn.poll().unwrap(), ClientEvent::ResponseEnd) {}
assert!(matches!(conn.poll().unwrap(), ClientEvent::Closed));
assert!(matches!(conn.send_body(b"x", &mut out), Ok(())));
assert!(matches!(conn.end_body(&mut out), Ok(())));
assert!(matches!(
conn.end_body(&mut out),
Err(HttpError::InvalidState(_))
));
}
#[test]
fn test_eof_before_response_is_an_error() {
let mut conn = ClientConn::new();
conn.send_request(&get(), &mut Vec::new()).unwrap();
conn.feed_eof();
assert_eq!(conn.poll().unwrap_err(), HttpError::UnexpectedEof);
}
}