From 2fac064c382d4953f19b031f68798ef52dbbab2c Mon Sep 17 00:00:00 2001 From: G2-Games Date: Thu, 24 Oct 2024 23:00:23 -0500 Subject: [PATCH] Properly use blake3 mmap --- Cargo.lock | 46 ++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 2 +- src/database.rs | 5 ----- src/main.rs | 36 +++++++++++++----------------------- web/request.js | 4 ++-- 5 files changed, 62 insertions(+), 31 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4e3d185..8320264 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -170,6 +170,8 @@ dependencies = [ "cc", "cfg-if", "constant_time_eq", + "memmap2", + "rayon-core", "serde", ] @@ -266,6 +268,31 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + [[package]] name = "darling" version = "0.20.10" @@ -790,6 +817,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "memmap2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" +dependencies = [ + "libc", +] + [[package]] name = "mime" version = "0.3.17" @@ -1055,6 +1091,16 @@ dependencies = [ "getrandom", ] +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.7" diff --git a/Cargo.toml b/Cargo.toml index 1fda834..c6c9828 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] bincode = { version = "2.0.0-rc.3", features = ["serde"] } -blake3 = { version = "1.5.4", features = ["serde"] } +blake3 = { version = "1.5.4", features = ["mmap", "rayon", "serde"] } chrono = { version = "0.4.38", features = ["serde"] } log = "0.4" maud = { version = "0.26", features = ["rocket"] } diff --git a/src/database.rs b/src/database.rs index b42eb09..10fb93e 100644 --- a/src/database.rs +++ b/src/database.rs @@ -57,9 +57,6 @@ pub struct MochiFile { /// The original name of the file name: String, - /// Size of the file in bytes - size: usize, - /// The location on disk (for deletion and management) filename: PathBuf, @@ -80,7 +77,6 @@ impl MochiFile { /// Create a new file that expires in `expiry`. pub fn new_with_expiry( name: &str, - size: usize, hash: Hash, filename: PathBuf, expire_duration: TimeDelta @@ -90,7 +86,6 @@ impl MochiFile { Self { name: name.to_string(), - size, filename, hash, upload_datetime: current, diff --git a/src/main.rs b/src/main.rs index 678e0f0..5106dd4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,7 +8,7 @@ use chrono::{DateTime, TimeDelta, Utc}; use database::{clean_loop, Database, MochiFile}; use log::info; use rocket::{ - data::{Limits, ToByteUnit}, form::Form, fs::{FileServer, Options, TempFile}, get, http::{ContentType, RawStr}, post, response::{content::{RawCss, RawJavaScript}, status::NotFound, Redirect}, routes, serde::{json::Json, Serialize}, tokio::{self, fs::File, io::AsyncReadExt}, Config, FromForm, State + data::{Limits, ToByteUnit}, form::Form, fs::{FileServer, Options, TempFile}, get, http::{ContentType, RawStr}, post, response::{content::{RawCss, RawJavaScript}, status::NotFound, Redirect}, routes, serde::{json::Json, Serialize}, tokio, Config, FromForm, State }; use settings::Settings; use strings::{parse_time_string, to_pretty_time}; @@ -155,14 +155,13 @@ async fn handle_upload( let filename = get_id( raw_name, - hash.0 + hash ); out_path.push(filename.clone()); let constructed_file = MochiFile::new_with_expiry( raw_name, - hash.1, - hash.0, + hash, out_path.clone(), expire_time ); @@ -175,8 +174,8 @@ async fn handle_upload( status: true, response: "File already exists", name: constructed_file.name().clone(), - url: "files/".to_string() + &filename, - hash: hash.0.to_hex()[0..10].to_string(), + url: filename, + hash: hash.to_hex()[0..10].to_string(), expires: Some(constructed_file.get_expiry()), ..Default::default() })) @@ -191,8 +190,8 @@ async fn handle_upload( Ok(Json(ClientResponse { status: true, name: constructed_file.name().clone(), - url: "files/".to_string() + &filename, - hash: hash.0.to_hex()[0..10].to_string(), + url: filename, + hash: hash.to_hex()[0..10].to_string(), expires: Some(constructed_file.get_expiry()), ..Default::default() })) @@ -207,11 +206,11 @@ struct ClientResponse { pub response: &'static str, - #[serde(skip_serializing_if = "String::is_empty")] + #[serde(skip_serializing_if = "str::is_empty")] pub name: String, - #[serde(skip_serializing_if = "String::is_empty")] + #[serde(skip_serializing_if = "str::is_empty")] pub url: String, - #[serde(skip_serializing_if = "String::is_empty")] + #[serde(skip_serializing_if = "str::is_empty")] pub hash: String, #[serde(skip_serializing_if = "Option::is_none")] pub expires: Option>, @@ -223,20 +222,11 @@ fn get_id(name: &str, hash: Hash) -> String { } /// Get the Blake3 hash of a file, without reading it all into memory, and also get the size -async fn hash_file>(input: &P) -> Result<(Hash, usize), std::io::Error> { - let mut file = File::open(input).await?; - let mut buf = vec![0; 5000000]; +async fn hash_file>(input: &P) -> Result { let mut hasher = blake3::Hasher::new(); + hasher.update_mmap_rayon(input)?; - let mut total = 0; - let mut bytes_read = None; - while bytes_read != Some(0) { - bytes_read = Some(file.read(&mut buf).await?); - total += bytes_read.unwrap(); - hasher.update(&buf[..bytes_read.unwrap()]); - } - - Ok((hasher.finalize(), total)) + Ok(hasher.finalize()) } /// An endpoint to obtain information about the server's capabilities diff --git a/web/request.js b/web/request.js index 80bf17f..e3d5dcf 100644 --- a/web/request.js +++ b/web/request.js @@ -57,7 +57,7 @@ function makeFinished(progressBar, progressText, linkRow, linkAddress, hash) { progressText.textContent = ""; const link = progressText.appendChild(document.createElement("a")); link.textContent = hash; - link.href = linkAddress; + link.href = "/files/" + linkAddress; link.target = "_blank"; let button = linkRow.appendChild(document.createElement("button")); @@ -68,7 +68,7 @@ function makeFinished(progressBar, progressText, linkRow, linkAddress, hash) { clearTimeout(buttonTimeout) } navigator.clipboard.writeText( - encodeURI(window.location.protocol + "//" + window.location.host + "/" + linkAddress) + encodeURI(window.location.protocol + "//" + window.location.host + "/files/" + linkAddress) ) button.textContent = "✅"; buttonTimeout = setTimeout(function() {