mirror of
https://codeberg.org/icewind/shelve.git
synced 2026-06-03 20:14:08 +02:00
use builtin multipart handling for form upload
This commit is contained in:
parent
c359239d27
commit
4d4108c83b
6 changed files with 163 additions and 712 deletions
|
|
@ -129,10 +129,6 @@ impl ExpireQueue {
|
|||
|
||||
expired
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.queue.lock().unwrap().len()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
111
src/main.rs
111
src/main.rs
|
|
@ -1,27 +1,26 @@
|
|||
use crate::expire_queue::{ExpireQueue, InvalidUploadIdError, UploadId};
|
||||
use crate::token::{UploadToken, ValidTokens};
|
||||
use dotenv::dotenv;
|
||||
use rocket::data::ToByteUnit;
|
||||
use rocket::fs::NamedFile;
|
||||
use futures_util::future::try_join_all;
|
||||
use rocket::data::{Limits, ToByteUnit};
|
||||
use rocket::form::Form;
|
||||
use rocket::fs::{FileName, NamedFile, TempFile};
|
||||
use rocket::request::FromParam;
|
||||
use rocket::response::Redirect;
|
||||
use rocket::{get, launch, post, put, routes, Data, Responder, State};
|
||||
use rocket::{get, launch, post, put, routes, Config, Data, FromForm, Responder, State};
|
||||
use rust_embed::RustEmbed;
|
||||
use serde::Serialize;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::env::{self, current_dir};
|
||||
use std::fs::{create_dir, read_dir, remove_dir_all, File};
|
||||
use std::fs::{create_dir, create_dir_all, read_dir, remove_dir_all};
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
use std::str::FromStr;
|
||||
use std::thread::{sleep, spawn, JoinHandle};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use upload::MultipartDatas;
|
||||
|
||||
mod expire_queue;
|
||||
mod token;
|
||||
mod upload;
|
||||
|
||||
impl<'r> FromParam<'r> for UploadId {
|
||||
type Error = InvalidUploadIdError;
|
||||
|
|
@ -81,36 +80,38 @@ enum UploadResponse {
|
|||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UploadData {
|
||||
struct UploadResponseData {
|
||||
success: bool,
|
||||
error: Option<&'static str>,
|
||||
error: Option<Cow<'static, str>>,
|
||||
urls: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(FromForm, Debug)]
|
||||
struct UploadData<'r> {
|
||||
expire: Option<u64>,
|
||||
ajax: bool,
|
||||
token: &'r str,
|
||||
files: Vec<TempFile<'r>>,
|
||||
}
|
||||
|
||||
#[post("/upload", data = "<data>")]
|
||||
fn post_upload(
|
||||
data: MultipartDatas,
|
||||
async fn post_upload(
|
||||
data: Form<UploadData<'_>>,
|
||||
accepted_tokens: &State<ValidTokens>,
|
||||
basedir: &State<PathBuf>,
|
||||
expire_queue: &State<ExpireQueue>,
|
||||
) -> UploadResponse {
|
||||
let mut fields: HashMap<String, String> = data
|
||||
.texts
|
||||
.into_iter()
|
||||
.map(|text| (text.key, text.value))
|
||||
.collect();
|
||||
let expire = fields
|
||||
.get("expire")
|
||||
.and_then(|expire| u64::from_str(expire).ok());
|
||||
let ajax = fields.get("ajax").is_some();
|
||||
let token = fields.remove("token").unwrap_or_default();
|
||||
let data = data.into_inner();
|
||||
let expire = data.expire;
|
||||
let ajax = data.ajax;
|
||||
let token = data.token;
|
||||
|
||||
if !accepted_tokens.contains(&token) {
|
||||
if !accepted_tokens.contains(token) {
|
||||
if ajax {
|
||||
UploadResponse::Data(
|
||||
serde_json::to_string(&UploadData {
|
||||
serde_json::to_string(&UploadResponseData {
|
||||
success: false,
|
||||
error: Some("invalid token"),
|
||||
error: Some("invalid token".into()),
|
||||
urls: None,
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
|
|
@ -119,28 +120,28 @@ fn post_upload(
|
|||
UploadResponse::Redirect(Redirect::to("/?error=invalid%20token"))
|
||||
}
|
||||
} else {
|
||||
match data
|
||||
.files
|
||||
.into_iter()
|
||||
.map(|file| {
|
||||
match try_join_all(data.files.into_iter().filter(|file| file.len() > 0).map(
|
||||
|mut file| async move {
|
||||
let id = UploadId::generate(now() + expire.unwrap_or(THOUSAND_YEARS));
|
||||
expire_queue.push(id);
|
||||
let name = &file.filename;
|
||||
let name = file.name().unwrap_or("upload");
|
||||
let ext = file.raw_name().and_then(filename_ext).unwrap_or_default();
|
||||
let name = format!("{}.{}", name, ext);
|
||||
let url = format!("{}/{}", id.as_string(), &name);
|
||||
|
||||
let mut path: PathBuf = basedir.join(id.as_string());
|
||||
create_dir(&path)?;
|
||||
path.push(name);
|
||||
|
||||
let mut file = File::open(&file.path)?;
|
||||
io::copy(&mut file, &mut File::create(&path)?)?;
|
||||
Ok(format!("{}/{}", id.as_string(), &name))
|
||||
})
|
||||
.collect::<io::Result<Vec<String>>>()
|
||||
file.persist_to(path).await.map(|_| url)
|
||||
},
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(urls) => {
|
||||
if ajax {
|
||||
UploadResponse::Data(
|
||||
serde_json::to_string(&UploadData {
|
||||
serde_json::to_string(&UploadResponseData {
|
||||
success: true,
|
||||
error: None,
|
||||
urls: Some(urls),
|
||||
|
|
@ -151,10 +152,10 @@ fn post_upload(
|
|||
UploadResponse::Redirect(Redirect::to(""))
|
||||
}
|
||||
}
|
||||
Err(_) => UploadResponse::Data(
|
||||
serde_json::to_string(&UploadData {
|
||||
Err(e) => UploadResponse::Data(
|
||||
serde_json::to_string(&UploadResponseData {
|
||||
success: false,
|
||||
error: Some("error while moving file"),
|
||||
error: Some(format!("error while moving file: {}", e).into()),
|
||||
urls: None,
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
|
|
@ -174,7 +175,7 @@ async fn download(id: UploadId, name: PathBuf, basedir: &State<PathBuf>) -> Opti
|
|||
}
|
||||
|
||||
#[launch]
|
||||
fn rockert() -> _ {
|
||||
fn rocket() -> _ {
|
||||
dotenv().ok();
|
||||
|
||||
let mut env: HashMap<_, _> = env::vars().collect();
|
||||
|
|
@ -191,9 +192,19 @@ fn rockert() -> _ {
|
|||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| current_dir().unwrap_or_default().join("data"));
|
||||
|
||||
let tmpdir = basedir.join("tmp");
|
||||
create_dir_all(&tmpdir).expect("failed to create tmp directory");
|
||||
|
||||
expire_job(basedir.clone(), expire_queue.clone());
|
||||
|
||||
rocket::build()
|
||||
let figment = Config::figment().merge(("temp_dir", tmpdir)).merge((
|
||||
"limits",
|
||||
Limits::new()
|
||||
.limit("file", 2.gibibytes())
|
||||
.limit("data-form", 2.gibibytes()),
|
||||
));
|
||||
|
||||
rocket::custom(figment)
|
||||
.manage(tokens)
|
||||
.manage(basedir.clone())
|
||||
.manage(expire_queue)
|
||||
|
|
@ -223,3 +234,23 @@ fn expire_job(expire_basedir: PathBuf, expire_queue: ExpireQueue) -> JoinHandle<
|
|||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn filename_ext(name: &FileName) -> Option<&str> {
|
||||
let raw = name.dangerous_unsafe_unsanitized_raw().as_str();
|
||||
let (name, ext) = raw.split_once('.')?;
|
||||
if name.len() > 0 && ext.len() < 8 && ext.chars().all(|c| c.is_ascii_alphanumeric() || c == '.')
|
||||
{
|
||||
Some(ext)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ext() {
|
||||
assert_eq!(Some("jpg"), filename_ext("foo.jpg".into()));
|
||||
assert_eq!(Some("tar.gz"), filename_ext("foo.tar.gz".into()));
|
||||
assert_eq!(None, filename_ext(".png".into()));
|
||||
assert_eq!(None, filename_ext("../foo.png".into()));
|
||||
assert_eq!(None, filename_ext("tmp/../foo.png".into()));
|
||||
}
|
||||
|
|
|
|||
158
src/upload.rs
158
src/upload.rs
|
|
@ -1,158 +0,0 @@
|
|||
use std::fs::{self, File};
|
||||
use std::io::{Cursor, Read, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use rocket::data::{self, FromData, Outcome, ToByteUnit};
|
||||
use rocket::{Data, Request};
|
||||
|
||||
use multipart::server::Multipart;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TextPart {
|
||||
pub key: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct FilePart {
|
||||
pub name: String,
|
||||
pub path: String,
|
||||
pub filename: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct MultipartDatas {
|
||||
pub texts: Vec<TextPart>,
|
||||
pub files: Vec<FilePart>,
|
||||
}
|
||||
|
||||
impl FilePart {
|
||||
pub fn persist(&self, p: &Path) {
|
||||
let s = Path::join(p, &self.filename);
|
||||
fs::copy(Path::new(&self.path), &s).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for FilePart {
|
||||
fn drop(&mut self) {
|
||||
fs::remove_file(Path::new(&self.path)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
const TMP_PATH: &str = "/tmp/rust_upload/";
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl<'r> FromData<'r> for MultipartDatas {
|
||||
type Error = String;
|
||||
|
||||
async fn from_data(
|
||||
request: &'r Request<'_>,
|
||||
data: Data<'r>,
|
||||
) -> data::Outcome<'r, Self, String> {
|
||||
let ct = request
|
||||
.headers()
|
||||
.get_one("Content-Type")
|
||||
.expect("no content-type");
|
||||
let idx = ct.find("boundary=").expect("no boundary");
|
||||
let boundary = &ct[(idx + "boundary=".len())..];
|
||||
|
||||
let mut d = Vec::new();
|
||||
data.open(2.gibibytes())
|
||||
.stream_to(&mut d)
|
||||
.await
|
||||
.expect("Unable to read");
|
||||
|
||||
let mut mp = Multipart::with_body(Cursor::new(d), boundary);
|
||||
let mut texts = Vec::new();
|
||||
let mut files = Vec::new();
|
||||
|
||||
let mut buffer = [0u8; 4096];
|
||||
|
||||
mp.foreach_entry(|entry| {
|
||||
let mut data = entry.data;
|
||||
if entry.headers.filename == None {
|
||||
let mut text_buffer = Vec::new();
|
||||
|
||||
loop {
|
||||
let c = match data.read(&mut buffer) {
|
||||
Ok(c) => c,
|
||||
Err(_err) => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if c == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
text_buffer.extend_from_slice(&buffer[..c]);
|
||||
}
|
||||
|
||||
let text = match String::from_utf8(text_buffer) {
|
||||
Ok(s) => s,
|
||||
Err(_err) => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
texts.push(TextPart {
|
||||
key: entry.headers.name.to_string(),
|
||||
value: text,
|
||||
});
|
||||
} else {
|
||||
let filename = entry.headers.filename.clone().unwrap();
|
||||
if !Path::new(TMP_PATH).exists() {
|
||||
fs::create_dir_all(TMP_PATH).unwrap();
|
||||
}
|
||||
|
||||
let target_path = Path::join(Path::new(TMP_PATH), &filename);
|
||||
|
||||
let mut file = match File::create(&target_path) {
|
||||
Ok(f) => f,
|
||||
Err(_err) => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut sum_c = 0u64;
|
||||
|
||||
loop {
|
||||
let c = match data.read(&mut buffer) {
|
||||
Ok(c) => c,
|
||||
Err(_err) => {
|
||||
try_delete(&target_path);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if c == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
sum_c = sum_c + c as u64;
|
||||
|
||||
match file.write(&buffer[..c]) {
|
||||
Ok(_) => (),
|
||||
Err(_err) => {
|
||||
try_delete(&target_path);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files.push(FilePart {
|
||||
name: entry.headers.name.to_string(),
|
||||
path: String::from(TMP_PATH) + &filename,
|
||||
filename: entry.headers.filename.clone().unwrap(),
|
||||
})
|
||||
}
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
Outcome::Success(MultipartDatas { texts, files })
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn try_delete<P: AsRef<Path>>(path: P) {
|
||||
if fs::remove_file(path.as_ref()).is_err() {}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue