use mime::Mime;
use super::field::{MultipartData, MultipartFile, ReadEntry, ReadEntryResult};
use self::SaveResult::*;
pub use tempdir::TempDir;
use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::path::{Path, PathBuf};
use std::{env, fs, io, mem};
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 SaveResult::Error(e),
}
)
);
#[must_use = "nothing saved to the filesystem yet"]
pub struct SaveBuilder<S> {
savable: S,
open_opts: OpenOptions,
size_limit: Option<u64>,
count_limit: Option<u32>,
}
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: savable,
open_opts: open_opts,
size_limit: None,
count_limit: None,
}
}
pub fn size_limit<L: Into<Option<u64>>>(mut self, limit: L) -> Self {
self.size_limit = limit.into();
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
}
}
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();
self
}
pub fn temp(self) -> EntriesSaveResult<M> {
self.temp_with_prefix("multipart-rs")
}
pub fn temp_with_prefix(self, prefix: &str) -> EntriesSaveResult<M> {
match TempDir::new(prefix) {
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.into())))
}
pub fn with_entries(mut self, mut entries: Entries) -> EntriesSaveResult<M> {
let mut count = 0;
loop {
let field = match ReadEntry::read_entry(self.savable) {
ReadEntryResult::Entry(field) => field,
ReadEntryResult::End(_) => break,
ReadEntryResult::Error(_, e) => return Partial (
PartialEntries {
entries: entries,
partial_file: None,
},
e.into(),
)
};
match field.data {
MultipartData::File(mut file) => {
match self.count_limit {
Some(limit) if count >= limit => return Partial (
PartialEntries {
entries: entries,
partial_file: Some(PartialFileField {
field_name: field.name,
source: file,
dest: None,
})
},
PartialReason::CountLimit,
),
_ => (),
}
count += 1;
match file.save().size_limit(self.size_limit).with_dir(&entries.save_dir) {
Full(saved_file) => {
self.savable = file.take_inner();
entries.mut_files_for(field.name).push(saved_file);
},
Partial(partial, reason) => return Partial(
PartialEntries {
entries: entries,
partial_file: Some(PartialFileField {
field_name: field.name,
source: file,
dest: Some(partial)
})
},
reason
),
Error(e) => return Partial(
PartialEntries {
entries: entries,
partial_file: Some(PartialFileField {
field_name: field.name,
source: file,
dest: None,
}),
},
e.into(),
),
}
},
MultipartData::Text(mut text) => {
self.savable = text.take_inner();
entries.fields.insert(field.name, text.text);
},
}
}
SaveResult::Full(entries)
}
}
impl<'m, M: 'm> SaveBuilder<&'m mut MultipartFile<M>> where MultipartFile<M>: BufRead {
pub fn temp(&mut self) -> FileSaveResult {
let path = env::temp_dir().join(rand_filename());
self.with_path(path)
}
pub fn with_filename(&mut self, filename: &str) -> FileSaveResult {
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) -> FileSaveResult {
let path = dir.as_ref().join(rand_filename());
self.with_path(path)
}
pub fn with_path<P: Into<PathBuf>>(&mut self, path: P) -> FileSaveResult {
let path = path.into();
let saved = SavedFile {
content_type: self.savable.content_type.clone(),
filename: self.savable.filename.clone(),
path: path,
size: 0,
};
let file = match create_dir_all(&saved.path).and_then(|_| self.open_opts.open(&saved.path)) {
Ok(file) => file,
Err(e) => return Partial(saved, e.into())
};
self.write_to(file).map(move |written| saved.with_size(written))
}
pub fn write_to<W: Write>(&mut self, mut dest: W) -> SaveResult<u64, u64> {
if let Some(limit) = self.size_limit {
let copied = match try_copy_buf(self.savable.take(limit), &mut dest) {
Full(copied) => copied,
other => return other,
};
match self.savable.fill_buf() {
Ok(buf) if buf.is_empty() => Full(copied),
Ok(_) => Partial(copied, PartialReason::SizeLimit),
Err(e) => Partial(copied, PartialReason::IoError(e))
}
} else {
try_copy_buf(&mut self.savable, &mut dest)
}
}
}
#[derive(Debug)]
pub struct SavedFile {
pub path: PathBuf,
pub filename: Option<String>,
pub content_type: Mime,
pub size: u64,
}
impl SavedFile {
fn with_size(self, size: u64) -> Self {
SavedFile { size: size, .. self }
}
}
#[derive(Debug)]
pub struct Entries {
pub fields: HashMap<String, String>,
pub files: HashMap<String, Vec<SavedFile>>,
pub save_dir: SaveDir,
}
impl Entries {
fn new(save_dir: SaveDir) -> Self {
Entries {
fields: HashMap::new(),
files: HashMap::new(),
save_dir: save_dir,
}
}
pub fn is_empty(&self) -> bool {
self.fields.is_empty() && self.files.is_empty()
}
fn mut_files_for(&mut self, field: String) -> &mut Vec<SavedFile> {
self.files.entry(field).or_insert_with(Vec::new)
}
}
#[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),
}
impl From<io::Error> for PartialReason {
fn from(e: io::Error) -> Self {
PartialReason::IoError(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 PartialFileField<M> {
pub field_name: String,
pub source: MultipartFile<M>,
pub dest: Option<SavedFile>,
}
#[derive(Debug)]
pub struct PartialEntries<M> {
pub entries: Entries,
pub partial_file: Option<PartialFileField<M>>,
}
impl<M> Into<Entries> for PartialEntries<M> {
fn into(self) -> Entries {
self.entries
}
}
impl<M> PartialEntries<M> {
pub fn keep_partial(mut self) -> Entries {
if let Some(partial_file) = self.partial_file {
if let Some(saved_file) = partial_file.dest {
self.entries.mut_files_for(partial_file.field_name).push(saved_file);
}
}
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 FileSaveResult = SaveResult<SavedFile, SavedFile>;
impl<M> 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, PartialReason::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_buf<R: BufRead, W: Write>(mut src: R, mut dest: W) -> 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; }
try_write_all(buf, &mut dest)
};
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>(mut buf: &[u8], mut dest: W) -> SaveResult<usize, usize> where W: Write {
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)
}