use std::borrow::Cow;
use std::fmt;
use std::ops::{Bound, RangeBounds};
use std::sync::{Arc, Mutex, PoisonError};
use bytes::{Bytes, BytesMut};
use futures_core::Stream;
use futures_util::{future, stream, StreamExt, TryStreamExt};
use crate::maybe::BoxStream;
use crate::{Body, BoxError, Error, MaybeSend, MaybeSync};
type Chunks = BoxStream<'static, Result<Bytes, BoxError>>;
type Once = Arc<Mutex<Option<Chunks>>>;
#[cfg(not(target_family = "wasm"))]
type Factory = Arc<dyn Fn() -> Chunks + Send + Sync>;
#[cfg(target_family = "wasm")]
type Factory = Arc<dyn Fn() -> Chunks>;
#[cfg_attr(
all(feature = "fs", not(target_family = "wasm")),
doc = "| [`Upload::path`], [`Upload::file`] | from the file's metadata | yes: the file is opened again |"
)]
#[cfg_attr(
not(all(feature = "fs", not(target_family = "wasm"))),
doc = "| `Upload::path`, `Upload::file` (feature `fs`, not on WebAssembly) | from the file's metadata | yes: the file is opened again |"
)]
#[cfg_attr(
all(feature = "fs", not(target_family = "wasm")),
doc = "| [`Upload::reader`], [`Upload::stream`] | [`Upload::len`] | no, unless [`Upload::buffer_for_retry`] holds it all |"
)]
#[cfg_attr(
not(all(feature = "fs", not(target_family = "wasm"))),
doc = "| `Upload::reader` (feature `fs`), [`Upload::stream`] | [`Upload::len`] | no, unless [`Upload::buffer_for_retry`] holds it all |"
)]
#[cfg_attr(
all(feature = "fs", not(target_family = "wasm")),
doc = "A path's length is read by [`Upload::resolve`], which generated clients call; [`Upload::file`] is"
)]
#[cfg_attr(
not(all(feature = "fs", not(target_family = "wasm"))),
doc = "A path's length is read by [`Upload::resolve`], which generated clients call; `Upload::file` is"
)]
#[derive(Clone)]
pub struct Upload {
source: Source,
pub(crate) file_name: Option<Cow<'static, str>>,
pub(crate) mime: Option<Cow<'static, str>>,
total: Option<u64>,
start: u64,
end: Option<u64>,
buffer_limit: Option<usize>,
}
#[derive(Clone)]
enum Source {
Bytes(Bytes),
#[cfg(feature = "fs")]
File(Arc<std::path::Path>),
Once(Once),
Factory(Factory),
}
impl Upload {
pub fn bytes(bytes: impl Into<Bytes>) -> Self {
let bytes = bytes.into();
let total = Some(bytes.len() as u64);
Self::new(Source::Bytes(bytes), total)
}
pub fn stream<S, B, E>(stream: S) -> Self
where
S: Stream<Item = Result<B, E>> + MaybeSend + 'static,
B: Into<Bytes> + 'static,
E: Into<BoxError> + 'static,
{
let chunks: Chunks = Box::pin(stream.map_ok(Into::into).map_err(Into::into));
Self::new(Source::Once(Arc::new(Mutex::new(Some(chunks)))), None)
}
pub fn from_fn<F, S, B, E>(make: F) -> Self
where
F: Fn() -> S + MaybeSend + MaybeSync + 'static,
S: Stream<Item = Result<B, E>> + MaybeSend + 'static,
B: Into<Bytes> + 'static,
E: Into<BoxError> + 'static,
{
let factory: Factory =
Arc::new(move || Box::pin(make().map_ok(Into::into).map_err(Into::into)));
Self::new(Source::Factory(factory), None)
}
fn new(source: Source, total: Option<u64>) -> Self {
Self {
source,
file_name: None,
mime: None,
total,
start: 0,
end: None,
buffer_limit: None,
}
}
pub fn file_name(mut self, file_name: impl Into<Cow<'static, str>>) -> Self {
self.file_name = Some(file_name.into());
self
}
pub fn mime(mut self, mime: impl Into<Cow<'static, str>>) -> Self {
self.mime = Some(mime.into());
self
}
pub fn name(&self) -> Option<&str> {
self.file_name.as_deref()
}
pub fn media_type(&self) -> Option<&str> {
self.mime.as_deref()
}
pub fn len(mut self, len: u64) -> Self {
self.total = Some(len);
self
}
pub fn slice(mut self, range: impl RangeBounds<u64>) -> Self {
let from = match range.start_bound() {
Bound::Included(&start) => start,
Bound::Excluded(&start) => start.saturating_add(1),
Bound::Unbounded => 0,
};
let to = match range.end_bound() {
Bound::Included(&end) => Some(end.saturating_add(1)),
Bound::Excluded(&end) => Some(end),
Bound::Unbounded => None,
};
let start = self.start.saturating_add(from);
let end = to.map(|to| self.start.saturating_add(to));
self.end = match (self.end, end) {
(Some(old), Some(new)) => Some(old.min(new)),
(old, new) => old.or(new),
};
self.start = match self.end {
Some(end) => start.min(end),
None => start,
};
self
}
pub fn buffer_for_retry(mut self, limit: usize) -> Self {
self.buffer_limit = Some(limit);
self
}
pub fn content_length(&self) -> Option<u64> {
self.total.map(|total| {
let end = self.end.unwrap_or(total).min(total);
end.saturating_sub(self.start.min(total))
})
}
pub fn is_replayable(&self) -> bool {
!matches!(self.source, Source::Once(_))
}
pub async fn resolve(mut self) -> Result<Self, Error> {
#[cfg(all(feature = "fs", target_family = "wasm"))]
if let Source::File(path) = &self.source {
return Err(Error::io(no_file_system(path)));
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
if let Source::File(path) = &self.source {
if self.total.is_none() {
let metadata = tokio::fs::metadata(&**path)
.await
.map_err(|error| Error::io(with_path(error, path)))?;
if !metadata.is_file() {
return Err(Error::io(std::io::Error::other(format!(
"{}: not a file",
path.display()
))));
}
self.total = Some(metadata.len());
}
return Ok(self);
}
let Some(limit) = self.buffer_limit else {
return Ok(self);
};
let Source::Once(once) = &self.source else {
return Ok(self);
};
let taken = once.lock().unwrap_or_else(PoisonError::into_inner).take();
let Some(mut chunks) = taken else {
return Ok(self);
};
let mut buffer = BytesMut::new();
while let Some(chunk) = chunks.next().await {
let chunk = chunk.map_err(Error::io)?;
if buffer.len() + chunk.len() > limit {
let read = stream::iter([Ok(buffer.freeze()), Ok(chunk)]).chain(chunks);
let chunks: Chunks = Box::pin(read);
self.source = Source::Once(Arc::new(Mutex::new(Some(chunks))));
return Ok(self);
}
buffer.extend_from_slice(&chunk);
}
self.total = Some(buffer.len() as u64);
self.source = Source::Bytes(buffer.freeze());
Ok(self)
}
}
#[cfg(feature = "fs")]
impl Upload {
#[cfg_attr(
not(target_family = "wasm"),
doc = "until the request is sent; [`Upload::resolve`] (or [`Upload::file`]) reads the length, so the body is sent with"
)]
#[cfg_attr(
target_family = "wasm",
doc = "until the request is sent; [`Upload::resolve`] (or `Upload::file`) reads the length, so the body is sent with"
)]
pub fn path(path: impl Into<std::path::PathBuf>) -> Self {
let path: Arc<std::path::Path> = Arc::from(path.into());
let file_name = path
.file_name()
.map(|name| Cow::Owned(name.to_string_lossy().into_owned()));
let mime = mime_guess::from_path(&path)
.first()
.map(|mime| Cow::Owned(mime.essence_str().to_owned()));
let mut upload = Self::new(Source::File(path), None);
upload.file_name = file_name;
upload.mime = mime;
upload
}
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
impl Upload {
pub async fn file(path: impl Into<std::path::PathBuf>) -> Result<Self, Error> {
Self::path(path).resolve().await
}
pub fn reader<R>(reader: R) -> Self
where
R: tokio::io::AsyncRead + Send + 'static,
{
Self::stream(read_chunks(Box::pin(reader), None))
}
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
fn read_chunks<R>(
reader: R,
expected: Option<u64>,
) -> impl Stream<Item = std::io::Result<Bytes>> + Send + 'static
where
R: tokio::io::AsyncRead + Send + Unpin + 'static,
{
use tokio::io::AsyncReadExt;
const CHUNK: usize = 64 * 1024;
stream::try_unfold((reader, 0u64), move |(mut reader, read)| async move {
let mut buffer = BytesMut::with_capacity(CHUNK);
let n = reader.read_buf(&mut buffer).await?;
if n == 0 {
return match expected {
Some(expected) if read < expected => Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
format!("the upload ended after {read} of its {expected} bytes"),
)),
_ => Ok(None),
};
}
Ok(Some((buffer.freeze(), (reader, read + n as u64))))
})
}
#[cfg(all(feature = "fs", target_family = "wasm"))]
fn no_file_system(path: &std::path::Path) -> std::io::Error {
std::io::Error::new(
std::io::ErrorKind::Unsupported,
format!(
"{}: no file system on WebAssembly: upload bytes or a stream instead",
path.display()
),
)
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
fn with_path(error: std::io::Error, path: &std::path::Path) -> std::io::Error {
std::io::Error::new(error.kind(), format!("{}: {error}", path.display()))
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
fn file_chunks(
path: Arc<std::path::Path>,
start: u64,
end: Option<u64>,
len: Option<u64>,
) -> Chunks {
use tokio::io::{AsyncReadExt, AsyncSeekExt};
let open = async move {
let mut file = tokio::fs::File::open(&*path)
.await
.map_err(|error| with_path(error, &path))?;
if start > 0 {
file.seek(std::io::SeekFrom::Start(start)).await?;
}
let limit = len
.or_else(|| end.map(|end| end.saturating_sub(start)))
.unwrap_or(u64::MAX);
Ok::<_, std::io::Error>(read_chunks(file.take(limit), len))
};
Box::pin(stream::once(open).try_flatten().map_err(BoxError::from))
}
fn range(chunks: Chunks, start: u64, end: Option<u64>) -> Chunks {
if start == 0 && end.is_none() {
return chunks;
}
Box::pin(stream::unfold(
(chunks, 0u64),
move |(mut chunks, mut position)| async move {
loop {
if end.is_some_and(|end| position >= end) {
return None;
}
let chunk = match chunks.next().await? {
Ok(chunk) => chunk,
Err(error) => return Some((Err(error), (chunks, position))),
};
let first = position;
position += chunk.len() as u64;
let from = start.max(first);
let to = end.map_or(position, |end| end.min(position));
if from < to {
let slice = chunk.slice((from - first) as usize..(to - first) as usize);
return Some((Ok(slice), (chunks, position)));
}
}
},
))
}
fn take_once(once: Once) -> Chunks {
let taken = stream::once(future::lazy(move |_| {
once.lock().unwrap_or_else(PoisonError::into_inner).take()
}));
Box::pin(taken.flat_map(|chunks| match chunks {
Some(chunks) => chunks,
None => Box::pin(stream::once(future::ready(Err(BoxError::from(
"this upload is a stream that was sent already: a stream is read once",
))))) as Chunks,
}))
}
impl From<Upload> for Body {
fn from(upload: Upload) -> Self {
let len = upload.content_length();
let (start, end) = (upload.start, upload.end);
match upload.source {
Source::Bytes(bytes) => {
let total = bytes.len() as u64;
let to = end.unwrap_or(total).min(total);
let from = start.min(to);
Body::from(bytes.slice(from as usize..to as usize))
}
#[cfg(all(feature = "fs", target_family = "wasm"))]
Source::File(path) => Body::wrap_stream(stream::once(future::ready(Err::<Bytes, _>(
no_file_system(&path),
)))),
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
Source::File(path) => Body::replay(
Arc::new(move || {
Body::sized_stream(file_chunks(path.clone(), start, end, len), len)
}),
len,
),
Source::Factory(make) => Body::replay(
Arc::new(move || Body::sized_stream(range(make(), start, end), len)),
len,
),
Source::Once(once) => Body::sized_stream(range(take_once(once), start, end), len),
}
}
}
impl From<Bytes> for Upload {
fn from(bytes: Bytes) -> Self {
Self::bytes(bytes)
}
}
impl From<Vec<u8>> for Upload {
fn from(bytes: Vec<u8>) -> Self {
Self::bytes(bytes)
}
}
impl From<&'static [u8]> for Upload {
fn from(bytes: &'static [u8]) -> Self {
Self::bytes(Bytes::from_static(bytes))
}
}
impl From<&Upload> for Upload {
fn from(upload: &Upload) -> Self {
upload.clone()
}
}
#[cfg(feature = "fs")]
impl From<std::path::PathBuf> for Upload {
fn from(path: std::path::PathBuf) -> Self {
Self::path(path)
}
}
#[cfg(feature = "fs")]
impl From<&std::path::PathBuf> for Upload {
fn from(path: &std::path::PathBuf) -> Self {
Self::path(path.clone())
}
}
#[cfg(feature = "fs")]
impl From<&std::path::Path> for Upload {
fn from(path: &std::path::Path) -> Self {
Self::path(path)
}
}
impl PartialEq for Upload {
fn eq(&self, other: &Self) -> bool {
let same_source = match (&self.source, &other.source) {
(Source::Bytes(a), Source::Bytes(b)) => a == b,
#[cfg(feature = "fs")]
(Source::File(a), Source::File(b)) => a == b,
(Source::Once(a), Source::Once(b)) => Arc::ptr_eq(a, b),
(Source::Factory(a), Source::Factory(b)) => {
std::ptr::addr_eq(Arc::as_ptr(a), Arc::as_ptr(b))
}
_ => false,
};
#[cfg(feature = "fs")]
let total = matches!(self.source, Source::File(_)) || self.total == other.total;
#[cfg(not(feature = "fs"))]
let total = self.total == other.total;
same_source
&& total
&& self.file_name == other.file_name
&& self.mime == other.mime
&& self.start == other.start
&& self.end == other.end
&& self.buffer_limit == other.buffer_limit
}
}
impl fmt::Debug for Upload {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut debug = f.debug_struct("Upload");
match &self.source {
Source::Bytes(bytes) => debug.field("bytes", &bytes.len()),
#[cfg(feature = "fs")]
Source::File(path) => debug.field("path", path),
Source::Once(_) => debug.field("stream", &true),
Source::Factory(_) => debug.field("from_fn", &true),
};
debug
.field("file_name", &self.file_name)
.field("mime", &self.mime)
.field("content_length", &self.content_length())
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
async fn read(upload: Upload) -> Bytes {
Body::from(upload).collect_bytes().await.unwrap()
}
fn chunks(
parts: &[&'static str],
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static {
stream::iter(
parts
.iter()
.map(|part| Ok(Bytes::from_static(part.as_bytes())))
.collect::<Vec<_>>(),
)
}
#[test]
fn uploads_are_equal_when_they_send_the_same_thing() {
let bytes = Upload::bytes(&b"abc"[..]).file_name("a.txt");
assert_eq!(bytes, bytes.clone());
assert_eq!(bytes, Upload::bytes(b"abc".to_vec()).file_name("a.txt"));
assert_ne!(bytes, Upload::bytes(&b"abc"[..]));
assert_ne!(bytes, bytes.clone().mime("text/plain"));
assert_ne!(bytes, bytes.clone().slice(1..));
let stream = Upload::stream(stream::iter([Ok::<_, std::io::Error>(&b"abc"[..])]));
assert_eq!(stream, stream.clone());
assert_ne!(
stream,
Upload::stream(stream::iter([Ok::<_, std::io::Error>(&b"abc"[..])]))
);
assert_ne!(stream, Upload::bytes(&b"abc"[..]));
let from_fn = Upload::from_fn(|| stream::iter([Ok::<_, std::io::Error>(&b"abc"[..])]));
assert_eq!(from_fn, from_fn.clone());
assert_ne!(from_fn, stream);
}
#[cfg(feature = "fs")]
#[tokio::test]
async fn paths_are_equal_as_paths() {
let file = tempfile::NamedTempFile::new().unwrap();
let path = Upload::path(file.path());
assert_eq!(path, Upload::from(file.path().to_path_buf()));
assert_eq!(path, path.clone().resolve().await.unwrap());
assert_ne!(path, Upload::path("other.txt"));
}
#[tokio::test]
async fn bytes_are_replayable_and_sliced_without_a_copy() {
let upload = Upload::bytes("hello world").slice(6..);
assert_eq!(upload.content_length(), Some(5));
let body = Body::from(upload);
assert_eq!(http_body::Body::size_hint(&body).exact(), Some(5));
let copy = body.try_clone().unwrap();
assert_eq!(body.collect_bytes().await.unwrap(), "world");
assert_eq!(copy.collect_bytes().await.unwrap(), "world");
assert_eq!(
read(Upload::bytes("hello").slice(1..=2).slice(1..)).await,
"l"
);
assert_eq!(read(Upload::bytes("hello").slice(10..)).await, "");
}
#[tokio::test]
async fn a_stream_is_read_once() {
let upload = Upload::stream(chunks(&["ab", "cd"]));
assert!(!upload.is_replayable());
let copy = upload.clone();
let body = Body::from(upload);
assert!(body.try_clone().is_none());
assert_eq!(body.collect_bytes().await.unwrap(), "abcd");
assert!(Body::from(copy).collect_bytes().await.is_err());
}
#[tokio::test]
async fn a_stream_is_sliced_across_chunks() {
assert_eq!(
read(Upload::stream(chunks(&["ab", "cd", "ef"])).slice(1..5)).await,
"bcde"
);
let sized = Upload::stream(chunks(&["ab", "cd"])).len(4).slice(1..);
assert_eq!(sized.content_length(), Some(3));
assert_eq!(
http_body::Body::size_hint(&Body::from(sized)).exact(),
Some(3)
);
}
#[tokio::test]
async fn from_fn_is_replayed() {
let upload = Upload::from_fn(|| chunks(&["x", "y"])).len(2);
assert!(upload.is_replayable());
let body = Body::from(upload);
let copy = body.try_clone().unwrap();
assert_eq!(body.collect_bytes().await.unwrap(), "xy");
assert_eq!(copy.collect_bytes().await.unwrap(), "xy");
}
#[tokio::test]
async fn buffer_for_retry_holds_a_short_stream() {
let short = Upload::stream(chunks(&["ab", "cd"]))
.buffer_for_retry(4)
.resolve()
.await
.unwrap();
assert!(short.is_replayable());
assert_eq!(short.content_length(), Some(4));
assert_eq!(read(short).await, "abcd");
let long = Upload::stream(chunks(&["ab", "cd", "e"]))
.buffer_for_retry(4)
.resolve()
.await
.unwrap();
assert!(!long.is_replayable());
assert_eq!(long.content_length(), None);
assert_eq!(read(long).await, "abcde");
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
#[tokio::test]
async fn a_file_is_resolved_reopened_and_sliced() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("notes.json");
std::fs::write(&path, "0123456789").unwrap();
let upload = Upload::path(&path);
assert_eq!(upload.file_name.as_deref(), Some("notes.json"));
assert_eq!(upload.mime.as_deref(), Some("application/json"));
assert_eq!(upload.content_length(), None);
let upload = upload.resolve().await.unwrap();
assert_eq!(upload.content_length(), Some(10));
let body = Body::from(upload.clone());
assert_eq!(http_body::Body::size_hint(&body).exact(), Some(10));
let copy = body.try_clone().unwrap();
assert_eq!(body.collect_bytes().await.unwrap(), "0123456789");
assert_eq!(copy.collect_bytes().await.unwrap(), "0123456789");
let slice = upload.slice(2..5);
assert_eq!(slice.content_length(), Some(3));
assert_eq!(read(slice.clone()).await, "234");
assert_eq!(read(slice.slice(1..)).await, "34");
let tail = Upload::file(&path).await.unwrap().slice(8..20);
assert_eq!(tail.content_length(), Some(2));
assert_eq!(read(tail).await, "89");
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
#[tokio::test]
async fn a_missing_file_fails_to_resolve() {
let error = Upload::file("/no/such/file.bin").await.unwrap_err();
assert!(error.is_io(), "{error:?}");
assert!(error.to_string().contains("/no/such/file.bin"), "{error}");
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
#[tokio::test]
async fn a_file_that_shrank_is_an_error() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("data.bin");
std::fs::write(&path, "0123456789").unwrap();
let upload = Upload::file(&path).await.unwrap();
std::fs::write(&path, "0123").unwrap();
assert!(Body::from(upload).collect_bytes().await.is_err());
}
#[cfg(all(feature = "fs", not(target_family = "wasm")))]
#[tokio::test]
async fn a_reader_is_read_once() {
let upload = Upload::reader(&b"from a reader"[..]);
assert!(!upload.is_replayable());
assert_eq!(read(upload).await, "from a reader");
}
}