From f16e7acadbf7367cf0a48ca69f419a70c1beddf6 Mon Sep 17 00:00:00 2001 From: kieran Date: Tue, 14 May 2024 10:58:21 +0100 Subject: [PATCH] Whitelist / better duplicate file response --- config.toml | 5 ++++- src/filesystem.rs | 18 ++++++++++++------ src/routes/blossom.rs | 20 ++++++++++++++++++-- src/routes/nip96.rs | 25 ++++++++++++++++++++++++- src/settings.rs | 3 +++ 5 files changed, 61 insertions(+), 10 deletions(-) diff --git a/config.toml b/config.toml index 8f76371..c5ded69 100644 --- a/config.toml +++ b/config.toml @@ -11,4 +11,7 @@ storage_dir = "./data" max_upload_bytes = 104857600 # Public facing url -public_url = "http://localhost:8000" \ No newline at end of file +public_url = "http://localhost:8000" + +# Whitelisted pubkeys, leave out to disable +# whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"] \ No newline at end of file diff --git a/src/filesystem.rs b/src/filesystem.rs index 614b72a..005f788 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -26,14 +26,14 @@ pub struct FileSystemResult { } pub struct FileStore { - path: String, + settings: Settings, processor: Arc>, } impl FileStore { pub fn new(settings: Settings) -> Self { Self { - path: settings.storage_dir, + settings, processor: Arc::new(Mutex::new(MediaProcessor::new())), } } @@ -50,6 +50,10 @@ impl FileStore { { let result = self.store_compress_file(stream, mime_type, compress).await?; let dst_path = self.map_path(&result.sha256); + if dst_path.exists() { + fs::remove_file(&result.path)?; + return Err(Error::msg("File already exists")); + } fs::create_dir_all(dst_path.parent().unwrap())?; if let Err(e) = fs::copy(&result.path, &dst_path) { fs::remove_file(&result.path)?; @@ -71,6 +75,7 @@ impl FileStore { let tmp_path = FileStore::map_temp(random_id); let mut file = File::options() .create(true) + .truncate(false) .write(true) .read(true) .open(tmp_path.clone()) @@ -83,7 +88,7 @@ impl FileStore { let start = SystemTime::now(); let proc_result = { let mut p_lock = self.processor.lock().expect("asd"); - p_lock.process_file(tmp_path.clone(), &mime_type)? + p_lock.process_file(tmp_path.clone(), mime_type)? }; if let FileProcessorResult::NewFile(new_temp) = proc_result { let old_size = tmp_path.metadata()?.len(); @@ -99,6 +104,7 @@ impl FileStore { fs::remove_file(tmp_path)?; file = File::options() .create(true) + .truncate(false) .write(true) .read(true) .open(new_temp.result.clone()) @@ -150,9 +156,9 @@ impl FileStore { fn map_path(&self, id: &Vec) -> PathBuf { let id = hex::encode(id); - Path::new(&self.path) - .join(id[0..2].to_string()) - .join(id[2..4].to_string()) + Path::new(&self.settings.storage_dir) + .join(&id[0..2]) + .join(&id[2..4]) .join(id) } } diff --git a/src/routes/blossom.rs b/src/routes/blossom.rs index 3e665b4..4a854ce 100644 --- a/src/routes/blossom.rs +++ b/src/routes/blossom.rs @@ -1,3 +1,5 @@ +use std::fs; + use chrono::Utc; use log::error; use nostr::prelude::hex; @@ -105,6 +107,12 @@ async fn upload( .content_type .unwrap_or("application/octet-stream".to_string()); + // check whitelist + if let Some(wl) = &settings.whitelist { + if !wl.contains(&auth.event.pubkey.to_hex()) { + return BlossomResponse::error("Not on whitelist"); + } + } match fs .put(data.open(ByteUnit::from(settings.max_upload_bytes)), &mime_type, false) .await @@ -126,7 +134,15 @@ async fn upload( created: Utc::now(), }; if let Err(e) = db.add_file(&f).await { - error!("{:?}", e); + error!("{}", e.to_string()); + let _ = fs::remove_file(blob.path); + if let Some(dbe) = e.as_database_error() { + if let Some(c) = dbe.code() { + if c == "23000" { + return BlossomResponse::error("File already exists"); + } + } + } BlossomResponse::error(format!("Error saving file (db): {}", e)) } else { BlossomResponse::BlobDescriptor(Json(BlobDescriptor::from_upload( @@ -136,7 +152,7 @@ async fn upload( } } Err(e) => { - error!("{:?}", e); + error!("{}", e.to_string()); BlossomResponse::error(format!("Error saving file (disk): {}", e)) } } diff --git a/src/routes/nip96.rs b/src/routes/nip96.rs index 293b270..d6168d0 100644 --- a/src/routes/nip96.rs +++ b/src/routes/nip96.rs @@ -1,6 +1,10 @@ +use std::borrow::Cow; use std::collections::HashMap; +use std::fs; use chrono::Utc; +use libc::remove; +use log::error; use rocket::{FromForm, Responder, Route, routes, State}; use rocket::form::Form; use rocket::fs::TempFile; @@ -158,6 +162,13 @@ async fn upload( }; let mime_type = form.media_type .unwrap_or("application/octet-stream"); + + // check whitelist + if let Some(wl) = &settings.whitelist { + if !wl.contains(&auth.event.pubkey.to_hex()) { + return Nip96Response::error("Not on whitelist"); + } + } match fs .put(file, mime_type, true) .await @@ -180,6 +191,15 @@ async fn upload( created: Utc::now(), }; if let Err(e) = db.add_file(&file_upload).await { + error!("{}", e.to_string()); + let _ = fs::remove_file(blob.path); + if let Some(dbe) = e.as_database_error() { + if let Some(c) = dbe.code() { + if c == "23000" { + return Nip96Response::error("File already exists"); + } + } + } return Nip96Response::error(&format!("Could not save file (db): {}", e)); } @@ -207,7 +227,10 @@ async fn upload( ..Default::default() })) } - Err(e) => return Nip96Response::error(&format!("Could not save file: {}", e)), + Err(e) => { + error!("{}", e.to_string()); + Nip96Response::error(&format!("Could not save file: {}", e)) + } } } diff --git a/src/settings.rs b/src/settings.rs index 0730ba9..e91240f 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -16,4 +16,7 @@ pub struct Settings { /// Public facing url pub public_url: String, + + /// Whitelisted pubkeys + pub whitelist: Option> }