303 lines
9.9 KiB
Rust
303 lines
9.9 KiB
Rust
use bytes::Bytes;
|
|
use dever_runtime::{
|
|
task::{Limits as TaskLimits, Runtime},
|
|
upload::{self, Limits},
|
|
};
|
|
use futures_util::stream::{self, StreamExt};
|
|
use http_body_util::StreamBody;
|
|
use hyper::body::Frame;
|
|
use std::{
|
|
convert::Infallible,
|
|
path::PathBuf,
|
|
sync::atomic::{AtomicUsize, Ordering},
|
|
time::Duration,
|
|
};
|
|
|
|
mod headers;
|
|
mod performance;
|
|
mod transport;
|
|
|
|
struct Directory(PathBuf);
|
|
|
|
#[derive(Debug, PartialEq, serde::Deserialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
struct Fields {
|
|
caption: String,
|
|
count: i64,
|
|
}
|
|
|
|
impl Directory {
|
|
fn new() -> Self {
|
|
static NEXT: AtomicUsize = AtomicUsize::new(0);
|
|
let path = std::env::temp_dir().join(format!(
|
|
"dever-upload-{}-{}",
|
|
std::process::id(),
|
|
NEXT.fetch_add(1, Ordering::Relaxed)
|
|
));
|
|
std::fs::create_dir(&path).unwrap();
|
|
Self(path)
|
|
}
|
|
|
|
fn temporary_count(&self) -> usize {
|
|
std::fs::read_dir(self.0.join("tmp")).map_or(0, |files| files.count())
|
|
}
|
|
}
|
|
|
|
impl Drop for Directory {
|
|
fn drop(&mut self) {
|
|
std::fs::remove_dir_all(&self.0).unwrap();
|
|
}
|
|
}
|
|
|
|
fn content(file: &[u8]) -> Vec<u8> {
|
|
let mut body = b"--boundary\r\nContent-Disposition: form-data; name=caption\r\n\r\nhello\r\n--boundary\r\nContent-Disposition: form-data; name=count\r\n\r\n7\r\n--boundary\r\nContent-Disposition: form-data; name=file; filename=image.bin\r\nContent-Type: application/octet-stream\r\n\r\n".to_vec();
|
|
body.extend_from_slice(file);
|
|
body.extend_from_slice(b"\r\n--boundary--\r\n");
|
|
body
|
|
}
|
|
|
|
fn body(
|
|
bytes: &[u8],
|
|
size: usize,
|
|
) -> impl hyper::body::Body<Data = Bytes, Error = Infallible> + Unpin + Send + 'static {
|
|
let chunks = bytes
|
|
.chunks(size)
|
|
.map(|chunk| Ok(Frame::data(Bytes::copy_from_slice(chunk))))
|
|
.collect::<Vec<_>>();
|
|
StreamBody::new(stream::iter(chunks))
|
|
}
|
|
|
|
#[test]
|
|
fn fragmented_files_preserve_false_boundaries_and_transfer_once() {
|
|
let directory = Directory::new();
|
|
let mut payload = vec![b'x'; 170_001];
|
|
payload.extend_from_slice(b"\r\n--boundaryX\r\n--boundary--X\r\n--boundary-");
|
|
let bytes = content(&payload);
|
|
Runtime::new(TaskLimits::default())
|
|
.unwrap()
|
|
.run(async |scope| {
|
|
for size in [1, 7, 16_384, 250_000] {
|
|
let mut decoded = upload::decode(
|
|
body(&bytes, size),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&Limits::default(),
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let file = decoded.take("file").unwrap();
|
|
assert_eq!(file.filename(), "image.bin");
|
|
assert_eq!(file.content_type(), "application/octet-stream");
|
|
assert_eq!(file.size(), payload.len() as i64);
|
|
assert_eq!(
|
|
decoded
|
|
.finish()
|
|
.unwrap()
|
|
.decode::<Fields>(dever_runtime::wire::decode)
|
|
.unwrap(),
|
|
Fields {
|
|
caption: "hello".into(),
|
|
count: 7
|
|
}
|
|
);
|
|
let key = upload::put(file, scope).await.unwrap();
|
|
assert_eq!(
|
|
std::fs::read(directory.0.join("upload").join(key.to_text())).unwrap(),
|
|
payload
|
|
);
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
}
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_or_limited_bodies_remove_every_temporary_file() {
|
|
let directory = Directory::new();
|
|
let bytes = content(&vec![b'x'; 20_000]);
|
|
Runtime::new(TaskLimits::default())
|
|
.unwrap()
|
|
.run(async |scope| {
|
|
let invalid = [
|
|
bytes[..bytes.len() - 8].to_vec(),
|
|
[bytes.clone(), b"epilogue".to_vec()].concat(),
|
|
String::from_utf8(bytes.clone())
|
|
.unwrap()
|
|
.replace("image.bin", "../escape")
|
|
.into_bytes(),
|
|
];
|
|
for bytes in invalid {
|
|
assert!(
|
|
upload::decode(
|
|
body(&bytes, 7),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&Limits::default(),
|
|
scope
|
|
)
|
|
.await
|
|
.is_err()
|
|
);
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
}
|
|
for limits in [
|
|
Limits {
|
|
total_bytes: 10_000,
|
|
file_bytes: 10_000,
|
|
..Limits::default()
|
|
},
|
|
Limits {
|
|
file_bytes: 10_000,
|
|
..Limits::default()
|
|
},
|
|
Limits {
|
|
field_bytes: 2,
|
|
..Limits::default()
|
|
},
|
|
Limits {
|
|
parts: 2,
|
|
fields: 2,
|
|
..Limits::default()
|
|
},
|
|
Limits {
|
|
header_bytes: 10,
|
|
..Limits::default()
|
|
},
|
|
] {
|
|
assert!(matches!(
|
|
upload::decode(
|
|
body(&bytes, 7),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&limits,
|
|
scope
|
|
)
|
|
.await,
|
|
Err(upload::Error::Limit)
|
|
));
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
}
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn duplicate_names_across_fields_and_files_remove_temporary_files() {
|
|
let directory = Directory::new();
|
|
Runtime::new(TaskLimits::default())
|
|
.unwrap()
|
|
.run(async |scope| {
|
|
for parts in [
|
|
[("same", false), ("same", false)],
|
|
[("same", false), ("same", true)],
|
|
[("same", true), ("same", false)],
|
|
[("same", true), ("same", true)],
|
|
[("first", true), ("second", true)],
|
|
] {
|
|
let mut source = String::new();
|
|
for (name, file) in parts {
|
|
let filename = if file { "; filename=example.txt" } else { "" };
|
|
source.push_str(&format!(
|
|
"--boundary\r\nContent-Disposition: form-data; name={name}{filename}\r\n\r\nvalue\r\n"
|
|
));
|
|
}
|
|
source.push_str("--boundary--\r\n");
|
|
assert!(matches!(
|
|
upload::decode(
|
|
body(source.as_bytes(), 7),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&Limits::default(),
|
|
scope,
|
|
)
|
|
.await,
|
|
Err(upload::Error::Invalid)
|
|
));
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
}
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn cancellation_while_waiting_for_more_file_bytes_removes_the_file() {
|
|
let directory = Directory::new();
|
|
let path = directory.0.clone();
|
|
let mut bytes = content(&vec![b'x'; 32_768]);
|
|
bytes.truncate(bytes.len() - 16);
|
|
Runtime::new(TaskLimits::default())
|
|
.unwrap()
|
|
.run(async |scope| {
|
|
let task = scope
|
|
.spawn(move |child| async move {
|
|
let frames = stream::once(std::future::ready(Ok::<_, Infallible>(
|
|
Frame::data(Bytes::from(bytes)),
|
|
)))
|
|
.chain(stream::pending());
|
|
upload::decode(
|
|
StreamBody::new(frames),
|
|
"multipart/form-data; boundary=boundary",
|
|
&path,
|
|
&Limits::default(),
|
|
&child,
|
|
)
|
|
.await
|
|
})
|
|
.unwrap();
|
|
tokio::time::timeout(Duration::from_secs(3), async {
|
|
while directory.temporary_count() == 0 {
|
|
tokio::task::yield_now().await;
|
|
}
|
|
})
|
|
.await
|
|
.unwrap();
|
|
task.cancel();
|
|
assert!(matches!(
|
|
task.wait().await,
|
|
Err(dever_runtime::task::Error::Cancelled) | Ok(Err(upload::Error::Cancelled))
|
|
));
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn dropping_unconsumed_file_and_failed_storage_remove_the_temporary() {
|
|
let directory = Directory::new();
|
|
let bytes = content(b"data");
|
|
Runtime::new(TaskLimits::default())
|
|
.unwrap()
|
|
.run(async |scope| {
|
|
let parsed = upload::decode(
|
|
body(&bytes, 7),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&Limits::default(),
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(directory.temporary_count(), 1);
|
|
drop(parsed);
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
std::fs::write(directory.0.join("upload"), b"not a directory").unwrap();
|
|
let mut parsed = upload::decode(
|
|
body(&bytes, 7),
|
|
"multipart/form-data; boundary=boundary",
|
|
&directory.0,
|
|
&Limits::default(),
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
upload::put(parsed.take("file").unwrap(), scope)
|
|
.await
|
|
.is_err()
|
|
);
|
|
assert_eq!(directory.temporary_count(), 0);
|
|
})
|
|
.unwrap();
|
|
}
|