use bytes::{Buf, BytesMut};
use tokio::io::{AsyncRead, AsyncReadExt};
use crate::helpers::fields::HeaderField;
use crate::helpers::sync::Timeout;
use crate::helpers::text::Text;
use crate::helpers::{hpack, qpack};
use crate::models::{Headers, Message, Method, URL, Version};
pub use crate::errors::Error;
impl From<hpack::Error> for Error {
fn from(err: hpack::Error) -> Self {
Self::Protocol(format!("hpack: {err}"))
}
}
impl From<qpack::Error> for Error {
fn from(err: qpack::Error) -> Self {
Self::Protocol(format!("qpack: {err}"))
}
}
#[derive(Default)]
pub struct StreamHasher(u64);
impl std::hash::Hasher for StreamHasher {
fn finish(&self) -> u64 {
self.0
}
fn write(&mut self, octets: &[u8]) {
for octet in octets {
self.write_u64(*octet as u64 ^ self.0);
}
}
fn write_u64(&mut self, value: u64) {
let mixed = value.wrapping_mul(0x9e37_79b9_7f4a_7c15);
self.0 = mixed ^ mixed >> 32;
}
fn write_u32(&mut self, value: u32) {
self.write_u64(value as u64);
}
}
pub type StreamMap<K, V> = std::collections::HashMap<K, V, std::hash::BuildHasherDefault<StreamHasher>>;
pub struct Buffer {
data: BytesMut,
chunk: usize,
chunk_size: usize,
eof: bool,
}
impl Buffer {
pub const DEFAULT_CHUNK_SIZE: usize = 16 * 1024;
pub const CHUNK_RAMP: usize = 8;
pub const MAXIMUM_CHUNK_SIZE: usize = 64 * 1024 * 1024;
pub fn oversized(capacity: usize, len: usize, idle_capacity: usize) -> bool {
capacity > idle_capacity && len <= idle_capacity / 2
}
pub fn reclaim_bytes(buffer: &mut BytesMut, idle_capacity: usize) {
if Self::oversized(buffer.capacity(), buffer.len(), idle_capacity) {
let mut fresh = BytesMut::new();
fresh.extend_from_slice(buffer);
*buffer = fresh;
}
}
pub fn reclaim_octets(buffer: &mut Vec<u8>, idle_capacity: usize) {
if Self::oversized(buffer.capacity(), buffer.len(), idle_capacity) {
buffer.shrink_to(idle_capacity / 2);
}
}
pub fn new() -> Self {
Self::with_chunk_size(Self::DEFAULT_CHUNK_SIZE)
}
pub fn with_chunk_size(chunk_size: usize) -> Self {
let mut buffer = Self { data: BytesMut::new(), chunk: 0, chunk_size: 0, eof: false };
buffer.set_chunk_size(chunk_size);
buffer
}
pub fn chunk_size(&self) -> usize {
self.chunk_size
}
pub fn set_chunk_size(&mut self, chunk_size: usize) {
self.chunk_size = chunk_size.clamp(1, Self::MAXIMUM_CHUNK_SIZE);
self.chunk = (self.chunk_size / Self::CHUNK_RAMP).max(1);
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn eof(&self) -> bool {
self.eof
}
pub fn as_slice(&self) -> &[u8] {
&self.data
}
pub fn as_bytes_mut(&mut self) -> &mut BytesMut {
&mut self.data
}
pub fn consume(&mut self, count: usize) {
self.data.advance(count.min(self.data.len()));
}
pub fn capacity(&self) -> usize {
self.data.capacity()
}
pub fn reclaim(&mut self, idle_capacity: usize) {
Self::reclaim_bytes(&mut self.data, idle_capacity);
}
pub fn take(&mut self, count: usize) -> BytesMut {
self.data.split_to(count.min(self.data.len()))
}
pub async fn fill<T>(&mut self, transport: &mut T, timeout: f64) -> Result<bool, Error>
where
T: AsyncRead + Unpin,
{
if self.eof {
return Ok(false);
}
self.data.reserve(self.chunk);
let read = Timeout::within(timeout, transport.read_buf(&mut self.data)).await??;
if read == 0 {
self.eof = true;
return Ok(false);
}
if read >= self.chunk {
self.chunk = self.chunk.saturating_mul(2).min(self.chunk_size);
}
Ok(true)
}
pub async fn require<T>(&mut self, transport: &mut T, count: usize, timeout: f64) -> Result<&[u8], Error>
where
T: AsyncRead + Unpin,
{
while self.data.len() < count && self.fill(transport, timeout).await? {}
if self.data.len() < count {
return Err(Error::Closed);
}
Ok(&self.data[..count])
}
}
impl Default for Buffer {
fn default() -> Self {
Self::new()
}
}
pub struct Fields;
impl Fields {
pub const PSEUDO_REQUEST: &[&str] = &[":method", ":scheme", ":authority", ":path", ":protocol"];
pub const PSEUDO_RESPONSE: &[&str] = &[":status"];
pub const CONNECTION_SPECIFIC: &[&str] = &["connection", "keep-alive", "proxy-connection", "transfer-encoding", "upgrade"];
pub const FORBIDDEN_TRAILERS: &[&str] = &["content-length", "expect", "host", "te", "trailer"];
pub fn pseudo(name: &str) -> bool {
matches!(name.as_bytes().first(), Some(b':'))
}
pub fn connection_specific(name: &str) -> bool {
let Some(first) = name.as_bytes().first() else {
return false;
};
match (name.len(), first) {
(7, b'u') => name == "upgrade",
(10, b'c') => name == "connection",
(10, b'k') => name == "keep-alive",
(16, b'p') => name == "proxy-connection",
(17, b't') => name == "transfer-encoding",
_ => false,
}
}
pub fn forbidden_trailer(name: &str) -> bool {
if Self::connection_specific(name) {
return true;
}
let Some(first) = name.as_bytes().first() else {
return false;
};
match (name.len(), first) {
(2, b't') => name == "te",
(4, b'h') => name == "host",
(6, b'e') => name == "expect",
(7, b't') => name == "trailer",
(14, b'c') => name == "content-length",
_ => false,
}
}
pub fn check(field: &HeaderField) -> Result<(), Error> {
if !HeaderField::is_lowercase_name(&field.name) {
return Err(Error::Protocol(format!("field name {:?} is not a lowercase token", field.name)));
}
if !HeaderField::is_value(&field.value) {
return Err(Error::Protocol(format!("field value of {:?} is not a field value", field.name)));
}
Ok(())
}
pub fn into_trailers(fields: Vec<HeaderField>) -> Result<Headers, Error> {
let mut present = 0u32;
for field in &fields {
if Self::pseudo(&field.name) {
return Err(Error::Protocol("trailer section carries a pseudo-header".into()));
}
Self::check(field)?;
if Self::forbidden_trailer(&field.name) {
return Err(Error::Protocol(format!("field {:?} cannot appear in a trailer section", field.name)));
}
present |= Headers::well_known(&field.name);
}
Ok(Headers::adopt(fields, present))
}
pub fn status(status_code: u16) -> Text {
if !(100..1000).contains(&status_code) {
return Text::from_string(status_code.to_string());
}
let digits = [
b'0' + (status_code / 100) as u8,
b'0' + (status_code / 10 % 10) as u8,
b'0' + (status_code % 10) as u8,
];
unsafe { Text::from_verified_ascii(&digits) }
}
pub fn of(message: &Message) -> Result<Vec<HeaderField>, Error> {
let mut fields = Vec::new();
Self::write(message, &mut fields)?;
Ok(fields)
}
pub fn write(message: &Message, fields: &mut Vec<HeaderField>) -> Result<(), Error> {
fields.reserve(Self::PSEUDO_REQUEST.len() + message.headers.as_ref().map_or(0, Headers::len));
if let Some(method) = message.method {
let target = message.target.as_deref().unwrap_or("/");
let headers = message.headers.as_ref();
fields.push(HeaderField::new(":method", method.as_str()));
let protocol = headers.and_then(|headers| headers.get(":protocol"));
if method == Method::CONNECT && protocol.is_none() {
if !URL::is_authority(target) {
return Err(Error::Protocol(format!("authority {target:?} is malformed")));
}
fields.push(HeaderField::new(":authority", target));
} else {
if !URL::is_target(target) {
return Err(Error::Protocol(format!("request target {target:?} is malformed")));
}
fields.push(HeaderField::new(":scheme", if message.security.secure { "https" } else { "http" }));
if let Some(authority) = headers.and_then(|headers| headers.get("host")) {
if !URL::is_authority(authority) {
return Err(Error::Protocol(format!("authority {authority:?} is malformed")));
}
fields.push(HeaderField::new(":authority", authority));
}
fields.push(HeaderField::new(":path", target));
if let Some(protocol) = protocol {
fields.push(HeaderField::new(":protocol", protocol));
}
}
} else if let Some(status_code) = message.status_code {
fields.push(HeaderField::new(":status", Self::status(status_code)));
} else {
return Err(Error::Protocol("message is neither a request nor a response".into()));
}
if let Some(headers) = &message.headers {
for field in headers.fields() {
if Self::pseudo(&field.name) || field.name == "host" {
continue;
}
Self::check(field)?;
if Self::connection_specific(&field.name) {
return Err(Error::Protocol(format!("connection-specific field {:?} cannot be framed", field.name)));
}
fields.push(field.clone());
}
}
Ok(())
}
pub fn message(fields: &[HeaderField], version: Version) -> Result<Message, Error> {
let mut owned = Vec::with_capacity(HeaderField::section_hint(fields.len()));
owned.extend_from_slice(fields);
Self::into_message(owned, version)
}
pub fn into_message(fields: Vec<HeaderField>, version: Version) -> Result<Message, Error> {
const PSEUDO_METHOD: u8 = 1 << 0;
const PSEUDO_STATUS: u8 = 1 << 1;
const PSEUDO_SCHEME: u8 = 1 << 2;
const PSEUDO_PATH: u8 = 1 << 3;
const PSEUDO_AUTHORITY: u8 = 1 << 4;
const PSEUDO_PROTOCOL: u8 = 1 << 5;
let mut message = Message::new(version);
let mut regular = false;
let mut present = 0u32;
let mut seen = 0u8;
let mut authority: Option<Text> = None;
let mut path: Option<Text> = None;
for field in &fields {
if !Self::pseudo(&field.name) {
Self::check(field)?;
if Self::connection_specific(&field.name) {
return Err(Error::Protocol(format!("connection-specific field {:?} cannot be framed", field.name)));
}
if field.name == "te" && field.value != "trailers" {
return Err(Error::Protocol("TE may only request trailers".into()));
}
present |= Headers::well_known(&field.name);
regular = true;
continue;
}
if regular {
return Err(Error::Protocol(format!("pseudo-header {:?} follows a regular field", field.name)));
}
if !HeaderField::is_value(&field.value) {
return Err(Error::Protocol(format!("pseudo-header {:?} carries a malformed value", field.name)));
}
let pseudo = match field.name.as_str() {
":method" => PSEUDO_METHOD,
":status" => PSEUDO_STATUS,
":scheme" => PSEUDO_SCHEME,
":path" => PSEUDO_PATH,
":authority" => PSEUDO_AUTHORITY,
":protocol" => PSEUDO_PROTOCOL,
name => return Err(Error::Protocol(format!("pseudo-header {name:?} is not defined"))),
};
if seen & pseudo != 0 {
return Err(Error::Protocol(format!("pseudo-header {:?} is repeated", field.name)));
}
seen |= pseudo;
match pseudo {
PSEUDO_METHOD => {
let method = field.value.parse().map_err(|_| Error::Protocol(format!("method {:?} is not recognised", field.value)))?;
message.method = Some(method);
}
PSEUDO_STATUS => {
if field.value.len() != 3 || !field.value.bytes().all(|byte| byte.is_ascii_digit()) {
return Err(Error::Protocol(format!("status {:?} is not three digits", field.value)));
}
let status_code = field.value.parse().map_err(|_| Error::Protocol(format!("status {:?} is not three digits", field.value)))?;
message.status_code = Some(status_code);
}
PSEUDO_SCHEME => message.security.secure = field.value == "https",
PSEUDO_AUTHORITY => {
if !URL::is_authority(&field.value) {
return Err(Error::Protocol(format!("authority {:?} is malformed", field.value)));
}
authority = Some(field.value.clone());
}
PSEUDO_PATH => {
if !URL::is_target(&field.value) {
return Err(Error::Protocol(format!("path {:?} is malformed", field.value)));
}
path = Some(field.value.clone());
}
_ => {}
}
}
let stray = if message.method.is_some() {
(seen & PSEUDO_STATUS != 0).then_some(":status")
} else {
[
(PSEUDO_METHOD, ":method"),
(PSEUDO_SCHEME, ":scheme"),
(PSEUDO_PATH, ":path"),
(PSEUDO_AUTHORITY, ":authority"),
(PSEUDO_PROTOCOL, ":protocol"),
]
.into_iter()
.find_map(|(bit, name)| (seen & bit != 0).then_some(name))
};
if let Some(name) = stray {
return Err(Error::Protocol(format!("pseudo-header {name:?} does not belong to this message")));
}
if message.method.is_none() && message.status_code.is_none() {
return Err(Error::Protocol("message has neither :method nor :status".into()));
}
let mut headers = Headers::adopt(fields, present);
if let Some(method) = message.method {
if method == Method::CONNECT && seen & PSEUDO_PROTOCOL == 0 {
message.target = authority.clone();
if message.target.is_none() || seen & (PSEUDO_SCHEME | PSEUDO_PATH) != 0 {
return Err(Error::Protocol("CONNECT carries an authority and nothing else".into()));
}
} else {
let path = path.unwrap_or_default();
if path.is_empty() || seen & PSEUDO_SCHEME == 0 {
return Err(Error::Protocol("request needs both a scheme and a non-empty path".into()));
}
message.target = Some(path);
}
if let Some(authority) = authority {
headers.insert("host", authority);
}
}
message.headers = Some(headers);
Ok(message)
}
}