261 lines
8.4 KiB
Rust
261 lines
8.4 KiB
Rust
use axum::{
|
|
body::Body,
|
|
extract::{Multipart, Path, State},
|
|
http::{header, HeaderValue, StatusCode},
|
|
response::{Html, IntoResponse, Response},
|
|
};
|
|
use std::path::PathBuf;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio_util::io::ReaderStream;
|
|
|
|
pub const MAX_TOTAL_STORAGE_BYTES: i64 = 10 * 1024 * 1024 * 1024;
|
|
pub const MAX_FILE_SIZE_BYTES: i64 = 10 * 1024 * 1024 * 1024;
|
|
|
|
use crate::web_app_state::{AppState, AppStateInner, AuthenticatedUserId};
|
|
|
|
pub struct UploadSuccess {
|
|
pub(crate) id: i32,
|
|
pub(crate) filename: String,
|
|
pub(crate) size: u64,
|
|
}
|
|
|
|
pub fn render_upload_page(state: &AppStateInner, error: Option<String>, success: Option<UploadSuccess>) -> Html<String> {
|
|
let mut ctx = tera::Context::new();
|
|
if let Some(error) = error {
|
|
ctx.insert("error", &error);
|
|
}
|
|
if let Some(success) = success {
|
|
ctx.insert("success_id", &success.id);
|
|
ctx.insert("success_filename", &success.filename);
|
|
ctx.insert("success_size", &success.size);
|
|
}
|
|
let body = state.tera
|
|
.render("upload.html", &ctx)
|
|
.expect("render upload");
|
|
Html(body)
|
|
}
|
|
|
|
pub async fn upload_get(
|
|
_: (),
|
|
_state: AppState,
|
|
_user: AuthenticatedUserId,) -> Html<String> {
|
|
render_upload_page(&_state, None, None)
|
|
}
|
|
|
|
pub async fn upload_post(
|
|
(headers, mut multipart): (axum::http::HeaderMap, Multipart),
|
|
_state: AppState,
|
|
_user: AuthenticatedUserId,
|
|
) -> Html<String> {
|
|
let content_length: u64 = match headers
|
|
.get(header::CONTENT_LENGTH)
|
|
.and_then(|value| value.to_str().ok())
|
|
.and_then(|value| value.parse::<u64>().ok())
|
|
{
|
|
Some(value) => value,
|
|
_ => return render_upload_page(&_state,
|
|
Some("Missing or invalid Content-Length".to_string()),
|
|
None,
|
|
)
|
|
};
|
|
|
|
if (content_length > MAX_FILE_SIZE_BYTES as u64){
|
|
return render_upload_page(&_state, Some("File size is too big".to_string()), None);
|
|
}
|
|
|
|
let old_total_size: i64 = match _state .db.query_one(
|
|
"SELECT COALESCE(SUM(size), 0)::BIGINT FROM public.stored_file",
|
|
&[],
|
|
).await {
|
|
Ok(row) => row.get::<_, i64>(0),
|
|
Err(err) => {
|
|
eprintln!("Database error: {err}");
|
|
return render_upload_page(&_state, Some("Server error".to_string()), None, );
|
|
}
|
|
};
|
|
|
|
assert!(0 <= old_total_size);
|
|
assert!(old_total_size <= MAX_TOTAL_STORAGE_BYTES);
|
|
|
|
let total_with_new: u64 = old_total_size as u64 + content_length;
|
|
|
|
if total_with_new > MAX_TOTAL_STORAGE_BYTES as u64 {
|
|
return render_upload_page(&_state,
|
|
Some("Storage limit exceeded (10 GiB)".to_string()), None,
|
|
);
|
|
}
|
|
|
|
let mut field = loop {
|
|
match multipart.next_field().await {
|
|
Ok(Some(field)) => {
|
|
if field.name() == Some("file") {
|
|
break field;
|
|
}
|
|
}
|
|
Ok(None) => {}
|
|
Err(err) =>
|
|
return render_upload_page(&_state,
|
|
Some(format!("Upload error: {err}")), None,
|
|
)
|
|
}
|
|
};
|
|
|
|
let filename = field
|
|
.file_name()
|
|
.map(|name| name.to_string())
|
|
.unwrap_or_else(|| "upload".to_string());
|
|
|
|
let row = match _state.db.query_one(
|
|
"INSERT INTO public.stored_file (filename, size) VALUES ($1, 0) RETURNING id",
|
|
&[&filename],
|
|
).await
|
|
{
|
|
Ok(row) => row,
|
|
Err(err) => {
|
|
return render_upload_page(&_state,
|
|
Some(format!("Database error: {err}")),
|
|
None,
|
|
);
|
|
}
|
|
};
|
|
let id: i32 = row.get::<_, i32>(0);
|
|
|
|
let path = PathBuf::from(&_state.config.file_storage).join(id.to_string());
|
|
let mut file = match tokio::fs::File::create(&path).await {
|
|
Ok(file) => file,
|
|
Err(err) => {
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
eprintln!("Cannot create file: {err}");
|
|
return render_upload_page(&_state,
|
|
Some("Filesystem error".to_string()), None,
|
|
);
|
|
}
|
|
};
|
|
|
|
if let Err(err) = file.set_len(content_length).await {
|
|
let _ = tokio::fs::remove_file(&path).await;
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
eprintln!("Cannot preallocate file: {err}");
|
|
return render_upload_page(&_state,
|
|
Some("Filesystem error".to_string()), None,
|
|
);
|
|
}
|
|
|
|
let mut written: u64 = 0;
|
|
loop {
|
|
match field.chunk().await {
|
|
Ok(Some(chunk)) => {
|
|
let chunk_len = chunk.len() as u64;
|
|
if written + chunk_len > content_length {
|
|
let _ = tokio::fs::remove_file(&path).await;
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
return render_upload_page(&_state,
|
|
Some("Your browser is bad".to_string()), None,
|
|
);
|
|
}
|
|
written += chunk_len;
|
|
{
|
|
let percent = 100. * (written as f32) / (content_length as f32);
|
|
println!("DEBUG: uploaded {percent}, cur chunk is {chunk_len}")
|
|
}
|
|
if let Err(err) = file.write_all(&chunk).await {
|
|
let _ = tokio::fs::remove_file(&path).await;
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
eprintln!("Can't write chunk to file: {err}");
|
|
return render_upload_page(&_state,
|
|
Some("Filesystem error".to_string()), None,
|
|
);
|
|
}
|
|
}
|
|
Ok(None) => break,
|
|
Err(err) => {
|
|
let _ = tokio::fs::remove_file(&path).await;
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
return render_upload_page(&_state,
|
|
Some(format!("Upload error: {err}")), None,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Err(err) = _state.db.execute(
|
|
"UPDATE public.stored_file SET size = $1 WHERE id = $2",
|
|
&[&(written as i64), &id],
|
|
).await {
|
|
let _ = tokio::fs::remove_file(&path).await;
|
|
let _ = _state.db.execute(
|
|
"DELETE FROM public.stored_file WHERE id = $1", &[&id]
|
|
).await;
|
|
|
|
return render_upload_page(&_state,
|
|
Some("Server error".to_string()), None,
|
|
);
|
|
}
|
|
|
|
render_upload_page(&_state,
|
|
None, Some(UploadSuccess {
|
|
id, filename, size: written, })
|
|
)
|
|
}
|
|
|
|
pub async fn get_file(
|
|
Path(id): Path<i32>,
|
|
state: AppState,
|
|
_user: AuthenticatedUserId,
|
|
) -> Response {
|
|
let row = match state
|
|
.db
|
|
.query_opt(
|
|
"SELECT filename, size FROM public.stored_file WHERE id = $1",
|
|
&[&id],
|
|
)
|
|
.await
|
|
{
|
|
Ok(Some(row)) => row,
|
|
Ok(None) => return StatusCode::NOT_FOUND.into_response(),
|
|
Err(err) => {
|
|
eprintln!("file lookup error: {err}");
|
|
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
|
|
}
|
|
};
|
|
|
|
let filename: String = row.get(0);
|
|
let size: i64 = row.get(1);
|
|
|
|
let path = PathBuf::from(&state.config.file_storage).join(id.to_string());
|
|
let file = match tokio::fs::File::open(&path).await {
|
|
Ok(file) => file,
|
|
Err(err) => {
|
|
eprintln!("file open error for {path:?}: {err}");
|
|
return StatusCode::NOT_FOUND.into_response();
|
|
}
|
|
};
|
|
|
|
let stream = ReaderStream::new(file);
|
|
let body = Body::from_stream(stream);
|
|
let mut response = Response::new(body);
|
|
if let Ok(value) = HeaderValue::from_str(&size.to_string()) {
|
|
response.headers_mut().insert(header::CONTENT_LENGTH, value);
|
|
}
|
|
response
|
|
.headers_mut()
|
|
.insert(header::CONTENT_TYPE, HeaderValue::from_static("application/octet-stream"));
|
|
response.headers_mut().insert(
|
|
header::CONTENT_DISPOSITION,
|
|
HeaderValue::from_str(&format!("attachment; filename=\"{}\"", filename)).unwrap_or_else(
|
|
|_| HeaderValue::from_static("attachment"),
|
|
),
|
|
);
|
|
response
|
|
}
|