tvl-depot/users/edef/turbofetch/src/lib.rs
Vincent Ambo bd8d6e824f chore(tvix/tools): move turbofetch to //users/edef
This was introduced in cl/9925 without any commit message, but this is clearly
not relevant to Tvix itself (it even says so in a comment in Cargo.toml).

Change-Id: I84f12d5145c3f53c9df23863f887bad913856c50
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12604
Tested-by: BuildkiteCI
Autosubmit: tazjin <tazjin@tvl.su>
Reviewed-by: edef <edef@edef.eu>
2024-10-12 23:02:50 +00:00

103 lines
3.1 KiB
Rust

use std::{mem::MaybeUninit, str};
use tokio::io::{self, AsyncRead, AsyncReadExt};
pub use buffer::Buffer;
mod buffer;
/// Read as much data into `buffer` as possible.
/// Returns [io::ErrorKind::OutOfMemory] if the buffer is already full.
async fn slurp(buffer: &mut Buffer, sock: &mut (impl AsyncRead + Unpin)) -> io::Result<()> {
match buffer.space() {
[] => Err(io::Error::new(io::ErrorKind::OutOfMemory, "buffer filled")),
buf => {
let n = sock.read(buf).await?;
if n == 0 {
return Err(io::ErrorKind::UnexpectedEof.into());
}
buffer.commit(n);
Ok(())
}
}
}
fn get_content_length(headers: &[httparse::Header]) -> io::Result<u64> {
for header in headers {
if header.name == "Transfer-Encoding" {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"Transfer-Encoding is unsupported",
));
}
if header.name == "Content-Length" {
return str::from_utf8(header.value)
.ok()
.and_then(|v| v.parse().ok())
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "invalid Content-Length")
});
}
}
Err(io::Error::new(
io::ErrorKind::InvalidData,
"Content-Length missing",
))
}
/// Read an HTTP response from `sock` using `buffer`, returning the response body.
/// Returns an error if anything but 200 OK is received.
///
/// The buffer must have enough space to contain the entire response body.
/// If there is not enough space, [io::ErrorKind::OutOfMemory] is returned.
///
/// The HTTP response must use `Content-Length`, without `Transfer-Encoding`.
pub async fn parse_response<'a>(
sock: &mut (impl AsyncRead + Unpin),
buffer: &'a mut Buffer,
) -> io::Result<&'a [u8]> {
let body_len = loop {
let mut headers = [MaybeUninit::uninit(); 16];
let mut response = httparse::Response::new(&mut []);
let status = httparse::ParserConfig::default()
.parse_response_with_uninit_headers(&mut response, buffer.data(), &mut headers)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
if let httparse::Status::Complete(n) = status {
buffer.consume(n);
let code = response.code.unwrap();
if code != 200 {
return Err(io::Error::new(
io::ErrorKind::Other,
format!("HTTP response {code}"),
));
}
break get_content_length(response.headers)?;
}
slurp(buffer, sock).await?;
};
let buf_len = buffer.space().len() + buffer.data().len();
if body_len > buf_len as u64 {
return Err(io::Error::new(
io::ErrorKind::OutOfMemory,
"HTTP response body does not fit in buffer",
));
}
let body_len = body_len as usize;
while buffer.data().len() < body_len {
slurp(buffer, sock).await?;
}
let data = buffer.data();
buffer.consume(body_len);
Ok(&data[..body_len])
}