check key field first

Signed-off-by: Amneesh Singh <natto@weirdnatto.in>
This commit is contained in:
2025-11-16 18:10:40 +05:30
parent 63a0268695
commit 1516480216

View File

@@ -227,79 +227,74 @@ fn clean_filename(filename: &str) -> String {
async fn upload(mut payload: Multipart) -> Result<impl IntoResponse, YamafError> { async fn upload(mut payload: Multipart) -> Result<impl IntoResponse, YamafError> {
let mut responses = Vec::new(); let mut responses = Vec::new();
let mut found_key = false;
if let Ok(ref key) = CONFIG.key {
if let Some(field) = payload.next_field().await.unwrap() {
if field.name() == Some("key") {
let bytes = field
.bytes()
.await
.map_err(|e| YamafError::BadRequest(format!("Error reading key: {e}")))?;
let s = String::from_utf8(bytes.to_vec())
.map_err(|_| YamafError::InternalError("Invalid key format".into()))?;
if s != *key {
return Err(YamafError::BadRequest("Wrong key".into()));
}
} else {
return Err(YamafError::BadRequest("Missing key".into()));
}
} else {
return Err(YamafError::BadRequest("Missing key".into()));
}
}
while let Some(mut field) = payload.next_field().await.unwrap() { while let Some(mut field) = payload.next_field().await.unwrap() {
match field.name() { if field.name() == Some("file") {
Some("key") => { let filename = field
if let Ok(ref key) = CONFIG.key { .file_name()
let bytes = field .map_or(format!("{}-upload", random(10)), |filename| {
.bytes() format!("{}-{}", random(4), clean_filename(filename))
.await });
.map_err(|e| YamafError::BadRequest(format!("Error reading key: {e}")))?;
let s = String::from_utf8(bytes.to_vec()) let save_path = std::path::Path::new(&CONFIG.root_dir).join(&filename);
.map_err(|_| YamafError::InternalError("Invalid key format".into()))?;
if s != *key { let mut file = fs::File::create(&save_path)
return Err(YamafError::BadRequest("Wrong key".into())); .await
} .map_err(|_| YamafError::InternalError("Internal i/o error".into()))?;
found_key = true; let mut written: usize = 0;
}
}
Some("file") => { while let Some(chunk) = field
if CONFIG.key.is_ok() && found_key == false { .chunk()
return Err(YamafError::BadRequest("Missing key".into())); .await
.map_err(|err| YamafError::InternalError(err.to_string()))?
{
use tokio::io::AsyncWriteExt;
written = written
.checked_add(chunk.len())
.ok_or_else(|| YamafError::BadRequest("File too large".into()))?;
if written > CONFIG.max_filesize {
_ = fs::remove_file(&save_path).await;
return Err(YamafError::FileTooBig(filename));
} }
let filename = field file.write_all(&chunk)
.file_name()
.map_or(format!("{}-upload", random(10)), |filename| {
format!("{}-{}", random(4), clean_filename(filename))
});
let save_path = std::path::Path::new(&CONFIG.root_dir).join(&filename);
let mut file = fs::File::create(&save_path)
.await .await
.map_err(|_| YamafError::InternalError("Internal i/o error".into()))?; .map_err(|_| YamafError::InternalError("Internal i/o error".into()))?;
}
let mut written: usize = 0; responses.push(format!(
while let Some(chunk) = field
.chunk()
.await
.map_err(|err| YamafError::InternalError(err.to_string()))?
{
use tokio::io::AsyncWriteExt;
written = written
.checked_add(chunk.len())
.ok_or_else(|| YamafError::BadRequest("File too large".into()))?;
if written > CONFIG.max_filesize {
_ = fs::remove_file(&save_path).await;
return Err(YamafError::FileTooBig(filename));
}
file.write_all(&chunk)
.await
.map_err(|_| YamafError::InternalError("Internal i/o error".into()))?;
}
responses.push(format!(
r#"<a href="{proto}://{host}/{file}">{proto}://{host}/{file}</a> (size ~ {size:.2}k)"#, r#"<a href="{proto}://{host}/{file}">{proto}://{host}/{file}</a> (size ~ {size:.2}k)"#,
proto = CONFIG.external_protocol, proto = CONFIG.external_protocol,
host = CONFIG.external_host, host = CONFIG.external_host,
file = filename, file = filename,
size = written as f64 / 1024 as f64 size = written as f64 / 1024 as f64
)); ));
}
None | Some(_) => {}
} }
} }