pub use server::buf_redux::BufReader;
pub use tempfile::TempDir;
use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::{self, File, OpenOptions};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::{cmp, env, io, mem, str, u32, u64};
use tempfile;
use server::field::{FieldHeaders, MultipartField, MultipartData, ReadEntry, ReadEntryResult};
use self::SaveResult::*;
use self::TextPolicy::*;
use self::PartialReason::*;
const RANDOM_FILENAME_LEN: usize = 12;
fn rand_filename() -> String {
::random_alphanumeric(RANDOM_FILENAME_LEN)
}
macro_rules! try_start (
($try:expr) => (
match $try {
Ok(val) => val,
Err(e) => return Error(e),
}
)
);
macro_rules! try_full (
($try:expr) => {
match $try {
Full(full) => full,
other => return other,
}
}
);
macro_rules! try_partial (
($try:expr) => {
match $try {
Full(full) => return Full(full.into()),
Partial(partial, reason) => (partial, reason),
Error(e) => return Error(e),
}
}
);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TextPolicy {
Try,
Force,
Ignore
}
#[must_use = "nothing saved to the filesystem yet"]
pub struct SaveBuilder<S> {
savable: S,
open_opts: OpenOptions,
size_limit: u64,
count_limit: u32,
memory_threshold: u64,
text_policy: TextPolicy,
}
impl<S> SaveBuilder<S> {
#[doc(hidden)]
pub fn new(savable: S) -> SaveBuilder<S> {
let mut open_opts = OpenOptions::new();
open_opts.write(true).create_new(true);
SaveBuilder {
savable,
open_opts,
size_limit: 8 * 1024 * 1024,
count_limit: 256,
memory_threshold: 10 * 1024,
text_policy: TextPolicy::Try,
}
}
pub fn size_limit<L: Into<Option<u64>>>(mut self, limit: L) -> Self {
self.size_limit = limit.into().unwrap_or(u64::MAX);
self
}
pub fn mod_open_opts<F: FnOnce(&mut OpenOptions)>(mut self, opts_fn: F) -> Self {
opts_fn(&mut self.open_opts);
self.open_opts.write(true);
self
}
pub fn memory_threshold(self, memory_threshold: u64) -> Self {
Self { memory_threshold, ..self }
}
pub fn try_text(self) -> Self {
Self { text_policy: TextPolicy::Try, ..self }
}
pub fn force_text(self) -> Self {
Self { text_policy: TextPolicy::Force, ..self}
}
pub fn ignore_text(self) -> Self {
Self { text_policy: TextPolicy::Ignore, ..self }
}
}
impl<M> SaveBuilder<M> where M: ReadEntry {
pub fn count_limit<L: Into<Option<u32>>>(mut self, count_limit: L) -> Self {
self.count_limit = count_limit.into().unwrap_or(u32::MAX);
self
}
pub fn temp(self) -> EntriesSaveResult<M> {
self.temp_with_prefix("multipart-rs")
}
pub fn temp_with_prefix(self, prefix: &str) -> EntriesSaveResult<M> {
match tempfile::Builder::new().prefix(prefix).tempdir() {
Ok(tempdir) => self.with_temp_dir(tempdir),
Err(e) => SaveResult::Error(e),
}
}
pub fn with_temp_dir(self, tempdir: TempDir) -> EntriesSaveResult<M> {
self.with_entries(Entries::new(SaveDir::Temp(tempdir)))
}
pub fn with_dir<P: Into<PathBuf>>(self, dir: P) -> EntriesSaveResult<M> {
let dir = dir.into();
try_start!(create_dir_all(&dir));
self.with_entries(Entries::new(SaveDir::Perm(dir)))
}
pub fn with_entries(self, mut entries: Entries) -> EntriesSaveResult<M> {
let SaveBuilder {
savable, open_opts, count_limit, size_limit,
memory_threshold, text_policy
} = self;
let mut res = ReadEntry::read_entry(savable);
let _ = entries.recount_fields();
let save_field = |field: &mut MultipartField<M>, entries: &Entries| {
let text_policy = if field.is_text() { text_policy } else { Ignore };
let mut saver = SaveBuilder {
savable: &mut field.data, open_opts: open_opts.clone(),
count_limit, size_limit, memory_threshold, text_policy
};
saver.with_dir(entries.save_dir.as_path())
};
while entries.fields_count < count_limit {
let mut field: MultipartField<M> = match res {
ReadEntryResult::Entry(field) => field,
ReadEntryResult::End(_) => return Full(entries), ReadEntryResult::Error(_, e) => return Partial (
PartialEntries {
entries,
partial: None,
},
e.into(),
)
};
let (dest, reason) = match save_field(&mut field, &entries) {
Full(saved) => {
entries.push_field(field.headers, saved);
res = ReadEntry::read_entry(field.data.into_inner());
continue;
},
Partial(saved, reason) => (Some(saved), reason),
Error(error) => (None, PartialReason::IoError(error)),
};
return Partial(
PartialEntries {
entries,
partial: Some(PartialSavedField {
source: field,
dest,
}),
},
reason
);
}
Partial(
PartialEntries {
entries,
partial: None,
},
PartialReason::CountLimit
)
}
}
impl<'m, M: 'm> SaveBuilder<&'m mut MultipartData<M>> where MultipartData<M>: BufRead {
pub fn temp(&mut self) -> FieldSaveResult {
let path = env::temp_dir().join(rand_filename());
self.with_path(path)
}
pub fn with_filename(&mut self, filename: &str) -> FieldSaveResult {
let mut tempdir = env::temp_dir();
tempdir.set_file_name(filename);
self.with_path(tempdir)
}
pub fn with_dir<P: AsRef<Path>>(&mut self, dir: P) -> FieldSaveResult {
let path = dir.as_ref().join(rand_filename());
self.with_path(path)
}
pub fn with_path<P: Into<PathBuf>>(&mut self, path: P) -> FieldSaveResult {
let bytes = if self.text_policy != Ignore {
let (text, reason) = try_partial!(self.save_text());
match reason {
SizeLimit if !self.cmp_size_limit(text.len()) => text.into_bytes(),
Utf8Error(_) if self.text_policy != Force => text.into_bytes(),
other => return Partial(text.into(), other),
}
} else {
Vec::new()
};
let (bytes, reason) = try_partial!(self.save_mem(bytes));
match reason {
SizeLimit if !self.cmp_size_limit(bytes.len()) => (),
other => return Partial(bytes.into(), other)
}
let path = path.into();
let mut file = match create_dir_all(&path).and_then(|_| self.open_opts.open(&path)) {
Ok(file) => file,
Err(e) => return Error(e),
};
let data = try_full!(
try_write_all(&bytes, &mut file)
.map(move |size| SavedData::File(path, size as u64))
);
self.write_to(file).map(move |written| data.add_size(written))
}
pub fn write_to<W: Write>(&mut self, mut dest: W) -> SaveResult<u64, u64> {
if self.size_limit < u64::MAX {
try_copy_limited(&mut self.savable, |buf| try_write_all(buf, &mut dest), self.size_limit)
} else {
try_read_buf(&mut self.savable, |buf| try_write_all(buf, &mut dest))
}
}
fn save_mem(&mut self, mut bytes: Vec<u8>) -> SaveResult<Vec<u8>, Vec<u8>> {
let pre_read = bytes.len() as u64;
match self.read_mem(|buf| { bytes.extend_from_slice(buf); Full(buf.len()) }, pre_read) {
Full(_) => Full(bytes),
Partial(_, reason) => Partial(bytes, reason),
Error(e) => if !bytes.is_empty() { Partial(bytes, e.into()) }
else { Error(e) }
}
}
fn save_text(&mut self) -> SaveResult<String, String> {
let mut string = String::new();
let res = self.read_mem(|buf| {
match str::from_utf8(buf) {
Ok(s) => { string.push_str(s); Full(buf.len()) },
Err(e) => if buf.len() < 4 {
Partial(0, e.into())
} else {
string.push_str(str::from_utf8(&buf[..e.valid_up_to()]).unwrap());
Full(e.valid_up_to())
}
}
}, 0);
match res {
Full(_) => Full(string),
Partial(_, reason) => Partial(string, reason),
Error(e) => Error(e),
}
}
fn read_mem<Wb: FnMut(&[u8]) -> SaveResult<usize, usize>>(&mut self, with_buf: Wb, pre_read: u64) -> SaveResult<u64, u64> {
let limit = cmp::min(self.size_limit, self.memory_threshold)
.saturating_sub(pre_read);
try_copy_limited(&mut self.savable, with_buf, limit)
}
fn cmp_size_limit(&self, size: usize) -> bool {
size as u64 >= self.size_limit
}
}
#[derive(Debug)]
pub struct SavedField {
pub headers: FieldHeaders,
pub data: SavedData,
}
#[derive(Debug)]
pub enum SavedData {
Text(String),
Bytes(Vec<u8>),
File(PathBuf, u64),
}
impl SavedData {
pub fn readable(&self) -> io::Result<DataReader> {
use self::SavedData::*;
match *self {
Text(ref text) => Ok(DataReader::Bytes(text.as_ref())),
Bytes(ref bytes) => Ok(DataReader::Bytes(bytes)),
File(ref path, _) => Ok(DataReader::File(BufReader::new(fs::File::open(path)?))),
}
}
pub fn size(&self) -> u64 {
use self::SavedData::*;
match *self {
Text(ref text) => text.len() as u64,
Bytes(ref bytes) => bytes.len() as u64,
File(_, size) => size,
}
}
pub fn is_memory(&self) -> bool {
use self::SavedData::*;
match *self {
Text(_) | Bytes(_) => true,
File(_, _) => false,
}
}
fn add_size(self, add: u64) -> Self {
use self::SavedData::File;
match self {
File(path, size) => File(path, size.saturating_add(add)),
other => other
}
}
}
impl From<String> for SavedData {
fn from(s: String) -> Self {
SavedData::Text(s)
}
}
impl From<Vec<u8>> for SavedData {
fn from(b: Vec<u8>) -> Self {
SavedData::Bytes(b)
}
}
pub enum DataReader<'a> {
Bytes(&'a [u8]),
File(BufReader<File>),
}
impl<'a> Read for DataReader<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
use self::DataReader::*;
match *self {
Bytes(ref mut bytes) => bytes.read(buf),
File(ref mut file) => file.read(buf),
}
}
}
impl<'a> BufRead for DataReader<'a> {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
use self::DataReader::*;
match *self {
Bytes(ref mut bytes) => bytes.fill_buf(),
File(ref mut file) => file.fill_buf(),
}
}
fn consume(&mut self, amt: usize) {
use self::DataReader::*;
match *self {
Bytes(ref mut bytes) => bytes.consume(amt),
File(ref mut file) => file.consume(amt),
}
}
}
#[derive(Debug)]
pub struct Entries {
pub fields: HashMap<Arc<str>, Vec<SavedField>>,
pub save_dir: SaveDir,
fields_count: u32,
}
impl Entries {
pub fn new(save_dir: SaveDir) -> Self {
Entries {
fields: HashMap::new(),
save_dir,
fields_count: 0,
}
}
pub fn is_empty(&self) -> bool {
self.fields.is_empty()
}
pub fn fields_count(&self) -> u32 {
self.fields_count
}
pub fn recount_fields(&mut self) -> u32 {
let fields_count = self.fields.values().map(Vec::len).sum();
self.fields_count = cmp::min(u32::MAX as usize, fields_count) as u32;
self.fields_count
}
fn push_field(&mut self, mut headers: FieldHeaders, data: SavedData) {
use std::collections::hash_map::Entry::*;
match self.fields.entry(headers.name.clone()) {
Vacant(vacant) => { vacant.insert(vec![SavedField { headers, data }]); },
Occupied(occupied) => {
headers.name = occupied.key().clone();
occupied.into_mut().push({ SavedField { headers, data }});
},
}
self.fields_count = self.fields_count.saturating_add(1);
}
pub fn print_debug(&self) -> io::Result<()> {
let stdout = io::stdout();
let stdout_lock = stdout.lock();
self.write_debug(stdout_lock)
}
pub fn write_debug<W: Write>(&self, mut writer: W) -> io::Result<()> {
for (name, entries) in &self.fields {
writeln!(writer, "Field {:?} has {} entries:", name, entries.len())?;
for (idx, field) in entries.iter().enumerate() {
let mut data = field.data.readable()?;
let headers = &field.headers;
writeln!(writer, "{}: {:?} ({:?}):", idx, headers.filename, headers.content_type)?;
io::copy(&mut data, &mut writer)?;
}
}
Ok(())
}
}
#[derive(Debug)]
pub enum SaveDir {
Temp(TempDir),
Perm(PathBuf),
}
impl SaveDir {
pub fn as_path(&self) -> &Path {
use self::SaveDir::*;
match *self {
Temp(ref tempdir) => tempdir.path(),
Perm(ref pathbuf) => &*pathbuf,
}
}
pub fn is_temporary(&self) -> bool {
use self::SaveDir::*;
match *self {
Temp(_) => true,
Perm(_) => false,
}
}
pub fn into_path(self) -> PathBuf {
use self::SaveDir::*;
match self {
Temp(tempdir) => tempdir.into_path(),
Perm(pathbuf) => pathbuf,
}
}
pub fn keep(&mut self) {
use self::SaveDir::*;
*self = match mem::replace(self, Perm(PathBuf::new())) {
Temp(tempdir) => Perm(tempdir.into_path()),
old_self => old_self,
};
}
pub fn delete(self) -> io::Result<()> {
use self::SaveDir::*;
match self {
Temp(tempdir) => tempdir.close(),
Perm(pathbuf) => fs::remove_dir_all(&pathbuf),
}
}
}
impl AsRef<Path> for SaveDir {
fn as_ref(&self) -> &Path {
self.as_path()
}
}
#[derive(Debug)]
pub enum PartialReason {
CountLimit,
SizeLimit,
IoError(io::Error),
Utf8Error(str::Utf8Error),
}
impl From<io::Error> for PartialReason {
fn from(e: io::Error) -> Self {
IoError(e)
}
}
impl From<str::Utf8Error> for PartialReason {
fn from(e: str::Utf8Error) -> Self {
Utf8Error(e)
}
}
impl PartialReason {
pub fn unwrap_err(self) -> io::Error {
self.expect_err("`PartialReason` was not `IoError`")
}
pub fn expect_err(self, msg: &str) -> io::Error {
match self {
PartialReason::IoError(e) => e,
_ => panic!("{}: {:?}", msg, self),
}
}
}
#[derive(Debug)]
pub struct PartialSavedField<M: ReadEntry> {
pub source: MultipartField<M>,
pub dest: Option<SavedData>,
}
#[derive(Debug)]
pub struct PartialEntries<M: ReadEntry> {
pub entries: Entries,
pub partial: Option<PartialSavedField<M>>,
}
impl<M: ReadEntry> Into<Entries> for PartialEntries<M> {
fn into(self) -> Entries {
self.entries
}
}
impl<M: ReadEntry> PartialEntries<M> {
pub fn keep_partial(mut self) -> Entries {
if let Some(partial) = self.partial {
if let Some(saved) = partial.dest {
self.entries.push_field(partial.source.headers, saved);
}
}
self.entries
}
}
#[derive(Debug)]
pub enum SaveResult<Success, Partial> {
Full(Success),
Partial(Partial, PartialReason),
Error(io::Error),
}
pub type EntriesSaveResult<M> = SaveResult<Entries, PartialEntries<M>>;
pub type FieldSaveResult = SaveResult<SavedData, SavedData>;
impl<M: ReadEntry> EntriesSaveResult<M> {
pub fn into_entries(self) -> Option<Entries> {
match self {
Full(entries) | Partial(PartialEntries { entries, .. }, _) => Some(entries),
Error(_) => None,
}
}
}
impl<S, P> SaveResult<S, P> where P: Into<S> {
pub fn okish(self) -> Option<S> {
self.into_opt_both().0
}
pub fn map<T, Map>(self, map: Map) -> SaveResult<T, T> where Map: FnOnce(S) -> T {
match self {
Full(full) => Full(map(full)),
Partial(partial, reason) => Partial(map(partial.into()), reason),
Error(e) => Error(e),
}
}
pub fn into_opt_both(self) -> (Option<S>, Option<io::Error>) {
match self {
Full(full) => (Some(full), None),
Partial(partial, IoError(e)) => (Some(partial.into()), Some(e)),
Partial(partial, _) => (Some(partial.into()), None),
Error(error) => (None, Some(error)),
}
}
pub fn into_result(self) -> io::Result<S> {
match self {
Full(entries) => Ok(entries),
Partial(partial, _) => Ok(partial.into()),
Error(error) => Err(error),
}
}
pub fn into_result_strict(self) -> io::Result<S> {
match self {
Full(entries) => Ok(entries),
Partial(_, PartialReason::IoError(e)) | Error(e) => Err(e),
Partial(partial, _) => Ok(partial.into()),
}
}
}
fn create_dir_all(path: &Path) -> io::Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
} else {
warn!("Attempting to save file in what looks like a root directory. File path: {:?}", path);
Ok(())
}
}
fn try_copy_limited<R: BufRead, Wb: FnMut(&[u8]) -> SaveResult<usize, usize>>(src: R, mut with_buf: Wb, limit: u64) -> SaveResult<u64, u64> {
let mut copied = 0u64;
try_read_buf(src, |buf| {
let new_copied = copied.saturating_add(buf.len() as u64);
if new_copied > limit { return Partial(0, PartialReason::SizeLimit) }
copied = new_copied;
with_buf(buf)
})
}
fn try_read_buf<R: BufRead, Wb: FnMut(&[u8]) -> SaveResult<usize, usize>>(mut src: R, mut with_buf: Wb) -> SaveResult<u64, u64> {
let mut total_copied = 0u64;
macro_rules! try_here (
($try:expr) => (
match $try {
Ok(val) => val,
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return if total_copied == 0 { Error(e) }
else { Partial(total_copied, e.into()) },
}
)
);
loop {
let res = {
let buf = try_here!(src.fill_buf());
if buf.is_empty() { break; }
with_buf(buf)
};
match res {
Full(copied) => { src.consume(copied); total_copied += copied as u64; }
Partial(copied, reason) => {
src.consume(copied); total_copied += copied as u64;
return Partial(total_copied, reason);
},
Error(err) => {
return Partial(total_copied, err.into());
}
}
}
Full(total_copied)
}
fn try_write_all<W: Write>(mut buf: &[u8], mut dest: W) -> SaveResult<usize, usize> {
let mut total_copied = 0;
macro_rules! try_here (
($try:expr) => (
match $try {
Ok(val) => val,
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return if total_copied == 0 { Error(e) }
else { Partial(total_copied, e.into()) },
}
)
);
while !buf.is_empty() {
match try_here!(dest.write(buf)) {
0 => try_here!(Err(io::Error::new(io::ErrorKind::WriteZero,
"failed to write whole buffer"))),
copied => {
buf = &buf[copied..];
total_copied += copied;
},
}
}
Full(total_copied)
}