diff options
author | Vincent Ambo <mail@tazj.in> | 2024-10-12T22·40+0300 |
---|---|---|
committer | clbot <clbot@tvl.fyi> | 2024-10-12T23·02+0000 |
commit | bd8d6e824f0bfc3c2755713af36f815d0ee82edb (patch) | |
tree | cf1b64998e2d199442bedfb25a85806e8b4e89ba /users/edef/turbofetch/src/lib.rs | |
parent | 398a9b531758297417280ccf06b82c85f3d236dd (diff) |
chore(tvix/tools): move turbofetch to //users/edef r/8799
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>
Diffstat (limited to 'users/edef/turbofetch/src/lib.rs')
-rw-r--r-- | users/edef/turbofetch/src/lib.rs | 103 |
1 files changed, 103 insertions, 0 deletions
diff --git a/users/edef/turbofetch/src/lib.rs b/users/edef/turbofetch/src/lib.rs new file mode 100644 index 000000000000..4b62fa4d75e7 --- /dev/null +++ b/users/edef/turbofetch/src/lib.rs @@ -0,0 +1,103 @@ +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]) +} |