extern crate rocket;
extern crate multipart;
extern crate chrono;
use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::path::{PathBuf, Path};
use std::sync::Arc;
use std::env;
use std::string;
use std::fs::{self, File};
use std::cmp::{Eq, PartialEq, PartialOrd, Ord, Ordering};
use chrono::prelude::*;
use rocket::Data;
use rocket::http::ContentType;
pub use rocket::http::hyper::mime;
use rocket::http::hyper::mime::{Mime, TopLevel, SubLevel};
use multipart::server::Multipart;
const BUFFER_SIZE: usize = 4096;
const DEFAULT_IN_MEMORY_DATA_LIMIT: u64 = 1 * 1024 * 1024;
const DEFAULT_FILE_DATA_LIMIT: u64 = 8 * 1024 * 1024;
#[derive(Debug)]
pub enum MultipartFormDataType {
Text,
Raw,
File,
}
#[derive(Debug)]
pub struct MultipartFormDataField<'a> {
pub t: MultipartFormDataType,
pub field_name: &'a str,
pub size_limit: u64,
pub content_type: Option<Vec<Mime>>,
}
impl<'a> MultipartFormDataField<'a> {
pub fn text(field_name: &'a str) -> MultipartFormDataField<'a> {
MultipartFormDataField {
t: MultipartFormDataType::Text,
field_name,
size_limit: DEFAULT_IN_MEMORY_DATA_LIMIT,
content_type: None,
}
}
pub fn bytes(field_name: &'a str) -> MultipartFormDataField<'a> {
Self::raw(field_name)
}
pub fn raw(field_name: &'a str) -> MultipartFormDataField<'a> {
MultipartFormDataField {
t: MultipartFormDataType::Raw,
field_name,
size_limit: DEFAULT_IN_MEMORY_DATA_LIMIT,
content_type: None,
}
}
pub fn file(field_name: &'a str) -> MultipartFormDataField<'a> {
MultipartFormDataField {
t: MultipartFormDataType::File,
field_name,
size_limit: DEFAULT_FILE_DATA_LIMIT,
content_type: None,
}
}
pub fn size_limit(mut self, size_limit: u64) -> MultipartFormDataField<'a> {
self.size_limit = size_limit;
self
}
pub fn content_type(mut self, content_type: Option<Mime>) -> MultipartFormDataField<'a> {
match content_type {
Some(content_type) => {
match self.content_type {
Some(mut v) => {
v.push(content_type);
self.content_type = Some(v);
}
None => {
self.content_type = Some(vec![content_type]);
}
}
}
None => self.content_type = None
}
self
}
pub fn content_type_by_string(mut self, content_type: Option<&'a str>) -> Result<MultipartFormDataField<'a>, ()> {
match content_type {
Some(content_type) => {
let content_type: Mime = content_type.parse()?;
match self.content_type {
Some(mut v) => {
v.push(content_type);
self.content_type = Some(v);
}
None => {
self.content_type = Some(vec![content_type]);
}
}
}
None => self.content_type = None
}
Ok(self)
}
}
impl<'a> PartialEq for MultipartFormDataField<'a> {
fn eq(&self, other: &Self) -> bool {
self.field_name.eq(other.field_name)
}
}
impl<'a> Eq for MultipartFormDataField<'a> {}
impl<'a> PartialOrd for MultipartFormDataField<'a> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
self.field_name.partial_cmp(other.field_name)
}
}
impl<'a> Ord for MultipartFormDataField<'a> {
fn cmp(&self, other: &Self) -> Ordering {
self.field_name.cmp(other.field_name)
}
}
#[derive(Debug)]
pub struct MultipartFormDataOptions<'a> {
pub temporary_dir: PathBuf,
pub allowed_fields: Vec<MultipartFormDataField<'a>>,
}
#[derive(Debug)]
pub struct MultipartFormData {
pub files: HashMap<Arc<String>, FileField>,
pub raw: HashMap<Arc<String>, RawField>,
pub texts: HashMap<Arc<String>, TextField>,
}
#[derive(Debug)]
pub struct SingleFileField {
pub content_type: Option<Mime>,
pub file_name: Option<String>,
pub path: PathBuf,
}
#[derive(Debug)]
pub enum FileField {
Single(SingleFileField),
Multiple(Vec<SingleFileField>),
}
#[derive(Debug)]
pub struct SingleRawField {
pub content_type: Option<Mime>,
pub file_name: Option<String>,
pub raw: Vec<u8>,
}
#[derive(Debug)]
pub enum RawField {
Single(SingleRawField),
Multiple(Vec<SingleRawField>),
}
#[derive(Debug)]
pub struct SingleTextField {
pub content_type: Option<Mime>,
pub file_name: Option<String>,
pub text: String,
}
#[derive(Debug)]
pub enum TextField {
Single(SingleTextField),
Multiple(Vec<SingleTextField>),
}
#[derive(Debug)]
pub enum MultipartFormDataError {
NotFormDataError,
BoundaryNotFoundError,
BodySizeTooLargeError,
IOError(io::Error),
FieldTypeAmbiguousError,
FromUtf8Error(string::FromUtf8Error),
DataTooLargeError(Arc<String>),
}
impl<'a> MultipartFormDataOptions<'a> {
pub fn new() -> MultipartFormDataOptions<'a> {
MultipartFormDataOptions {
temporary_dir: env::temp_dir(),
allowed_fields: Vec::new(),
}
}
}
impl MultipartFormData {
pub fn parse(content_type: &ContentType, data: Data, mut options: MultipartFormDataOptions) -> Result<MultipartFormData, MultipartFormDataError> {
if !content_type.is_form_data() {
return Err(MultipartFormDataError::NotFormDataError);
}
let (_, boundary) = match content_type.params().find(|&(k, _)| k == "boundary") {
Some(s) => s,
None => return Err(MultipartFormDataError::BoundaryNotFoundError)
};
options.allowed_fields.sort();
let mut multipart = Multipart::with_body(data.open(), boundary);
let mut files = HashMap::new();
let mut raw = HashMap::new();
let mut texts = HashMap::new();
if !files.is_empty() {
let path = options.temporary_dir.as_path();
if path.exists() {
if !path.is_dir() {
return Err(MultipartFormDataError::IOError(io::Error::new(io::ErrorKind::AlreadyExists, "the temporary path exists and it is not a directory")));
}
} else {
fs::create_dir_all(path).map_err(|err| MultipartFormDataError::IOError(err))?;
}
}
loop {
match multipart.read_entry().map_err(|err| MultipartFormDataError::IOError(err))? {
Some(entry) => {
let field_name = entry.headers.name;
let content_type = entry.headers.content_type;
'accept: loop {
if let Ok(vi) = options.allowed_fields.binary_search_by(|f| f.field_name.cmp(&field_name.as_str())) {
{
let field_ref = &options.allowed_fields[vi];
if let Some(content_type_ref) = &field_ref.content_type { let mut mat = false;
let (top, sub) = match &content_type {
Some(content_type) => {
let Mime(top, sub, _) = content_type;
(Some(top), Some(sub))
}
None => (None, None)
};
for content_type_ref in content_type_ref {
let Mime(top_ref, sub_ref, _) = content_type_ref;
if top_ref.ne(&TopLevel::Star) {
if let Some(top) = top {
if top_ref.ne(top) {
continue;
}
} else {
continue;
}
}
if sub_ref.ne(&SubLevel::Star) {
if let Some(sub) = sub {
if sub_ref.ne(sub) {
continue;
}
} else {
continue;
}
}
mat = true;
break;
}
if !mat {
continue 'accept;
}
}
}
let field = options.allowed_fields.remove(vi);
let mut data = entry.data;
let mut buffer = [0u8; BUFFER_SIZE];
match field.t {
MultipartFormDataType::File => {
let now = Utc::now();
let target_file_name = format!("rs-{}", now.timestamp_nanos());
let target_path = {
let mut i = 0usize;
let mut p;
loop {
p = if i == 0 {
Path::join(&options.temporary_dir, &target_file_name)
} else {
Path::join(&options.temporary_dir, format!("{}-{}", &target_file_name, i))
};
if !p.exists() {
break;
}
i += 1;
}
p
};
let mut file = File::create(&target_path).map_err(|err| MultipartFormDataError::IOError(err))?;
let mut sum_c = 0u64;
loop {
let c = data.read(&mut buffer).map_err(|err| {
try_delete(&target_path);
MultipartFormDataError::IOError(err)
})?;
if c == 0 {
break;
}
sum_c += c as u64;
if sum_c > field.size_limit {
try_delete(&target_path);
return Err(MultipartFormDataError::DataTooLargeError(field_name));
}
file.write(&buffer[..c]).map_err(|err| {
try_delete(&target_path);
MultipartFormDataError::IOError(err)
})?;
}
let file_name = entry.headers.filename;
let f = SingleFileField {
content_type,
file_name,
path: target_path,
};
match files.remove(&field_name) {
Some(field) => {
match field {
FileField::Single(t) => {
let v = vec![t, f];
files.insert(field_name, FileField::Multiple(v));
}
FileField::Multiple(mut v) => {
v.push(f);
files.insert(field_name, FileField::Multiple(v));
}
}
}
None => {
files.insert(field_name, FileField::Single(f));
}
}
}
MultipartFormDataType::Raw => {
let mut bytes = Vec::new();
loop {
let c = data.read(&mut buffer).map_err(|err| MultipartFormDataError::IOError(err))?;
if c == 0 {
break;
}
if bytes.len() as u64 + c as u64 > field.size_limit {
return Err(MultipartFormDataError::DataTooLargeError(field_name));
}
bytes.extend_from_slice(&buffer[..c]);
}
let file_name = entry.headers.filename;
let f = SingleRawField {
content_type,
file_name,
raw: bytes,
};
match raw.remove(&field_name) {
Some(field) => {
match field {
RawField::Single(t) => {
let v = vec![t, f];
raw.insert(field_name, RawField::Multiple(v));
}
RawField::Multiple(mut v) => {
v.push(f);
raw.insert(field_name, RawField::Multiple(v));
}
}
}
None => {
raw.insert(field_name, RawField::Single(f));
}
}
}
MultipartFormDataType::Text => {
let mut text_buffer = Vec::new();
loop {
let c = data.read(&mut buffer).map_err(|err| MultipartFormDataError::IOError(err))?;
if c == 0 {
break;
}
if text_buffer.len() as u64 + c as u64 > field.size_limit {
return Err(MultipartFormDataError::DataTooLargeError(field_name));
}
text_buffer.extend_from_slice(&buffer[..c]);
}
let text = String::from_utf8(text_buffer).map_err(|err| MultipartFormDataError::FromUtf8Error(err))?;
let file_name = entry.headers.filename;
let f = SingleTextField {
content_type,
file_name,
text,
};
match texts.remove(&field_name) {
Some(field) => {
match field {
TextField::Single(t) => {
let v = vec![t, f];
texts.insert(field_name, TextField::Multiple(v));
}
TextField::Multiple(mut v) => {
v.push(f);
texts.insert(field_name, TextField::Multiple(v));
}
}
}
None => {
texts.insert(field_name, TextField::Single(f));
}
}
}
}
break 'accept;
} else {
break 'accept;
}
}
}
None => {
break;
}
}
}
Ok(MultipartFormData {
files,
raw,
texts,
})
}
}
impl Drop for MultipartFormData {
fn drop(&mut self) {
let files = &self.files;
for (_, field) in files {
match field {
FileField::Single(f) => {
try_delete(&f.path);
}
FileField::Multiple(vf) => {
for f in vf {
try_delete(&f.path);
}
}
}
}
}
}
fn try_delete<P: AsRef<Path>>(path: P) {
if let Err(_) = fs::remove_file(path.as_ref()) {}
}