Merge pull request #313 from NordSecurity/FILE-579-checksum-progress-granularity

dfetti · LukasPukenis · commit cffea989dd7e · 2024-04-10T10:37:11.000+03:00
Add checksum_events_granularity config
diff --git a/changelog.md b/changelog.md
@@ -33,6 +33,7 @@
 * Disallow downloading file for which any path component is larger than 250 characters
 * Fix ocassional missing of `TransferPaused` event when toggling libdrop on and off quickly
 * Report file transfer error in case file subpath contains perent directory `..`
+* Add checksum_events_granularity config
 
 ---
 <br>
diff --git a/drop-config/src/lib.rs b/drop-config/src/lib.rs
@@ -13,6 +13,9 @@ pub struct DropConfig {
     pub storage_path: String,
     // If set the checksum events will be emited for every file of this or bigger size
     pub checksum_events_size_threshold: Option<usize>,
+    // If set the checksum events will be emited for every checksum_events_granularity bytes
+    // Default value is 256KB.
+    pub checksum_events_granularity: u64,
     pub connection_retries: u32,
 }
 
@@ -23,6 +26,7 @@ impl Default for DropConfig {
             transfer_file_limit: 1000,
             storage_path: "libdrop.sqlite".to_string(),
             checksum_events_size_threshold: None,
+            checksum_events_granularity: 256 * 1024,
             connection_retries: 5,
         }
     }
diff --git a/drop-transfer/src/file/mod.rs b/drop-transfer/src/file/mod.rs
@@ -299,13 +299,14 @@ impl FileToSend {
         &self,
         limit: u64,
         progress_cb: Option<F>,
+        event_granularity: Option<u64>,
     ) -> crate::Result<[u8; 32]>
     where
         F: FnMut(u64) -> Fut + Send + Sync,
         Fut: Future<Output = ()>,
     {
         let reader = reader::open(&self.source)?.take(limit);
-        let csum = checksum(reader, progress_cb).await?;
+        let csum = checksum(reader, progress_cb, event_granularity).await?;
         Ok(csum)
     }
 }
@@ -315,6 +316,7 @@ impl FileToSend {
 pub async fn checksum<F, Fut>(
     reader: impl io::Read,
     mut progress_cb: Option<F>,
+    event_granularity: Option<u64>,
 ) -> io::Result<[u8; 32]>
 where
     F: FnMut(u64) -> Fut + Send + Sync,
@@ -325,6 +327,7 @@ where
     let mut reader = io::BufReader::with_capacity(CHECKSUM_CHUNK_SIZE, reader);
 
     let mut total_n: u64 = 0;
+    let mut from_last_event: u64 = 0;
     loop {
         let buf = reader.fill_buf()?;
         if buf.is_empty() {
@@ -337,9 +340,15 @@ where
         reader.consume(n);
 
         total_n += n as u64;
-        if let Some(progress_cb) = progress_cb.as_mut() {
-            progress_cb(total_n).await;
+
+        if let (Some(progress_cb), Some(granularity)) = (progress_cb.as_mut(), event_granularity) {
+            from_last_event += n as u64;
+            while from_last_event >= granularity {
+                progress_cb(total_n).await;
+                from_last_event -= granularity;
+            }
         }
+
         // Since these are all blocking operation we need to give tokio runtime a
         // timeslice
         tokio::task::yield_now().await;
@@ -375,6 +384,7 @@ mod tests {
         let csum = super::checksum(
             &mut &TEST[..],
             None::<fn(u64) -> futures::future::Ready<()>>,
+            None,
         )
         .await
         .unwrap();
@@ -391,7 +401,7 @@ mod tests {
 
             let size = TEST.len() as _;
             let file = super::FileToSend::from_path(tmp.path(), size).unwrap();
-            file.checksum(size, None::<fn(u64) -> futures::future::Ready<()>>)
+            file.checksum(size, None::<fn(u64) -> futures::future::Ready<()>>, None)
                 .await
                 .unwrap()
         };
@@ -421,8 +431,11 @@ mod tests {
         let mut cx = Context::from_waker(&waker);
 
         let mut cursor = io::Cursor::new(&buf);
-        let mut future =
-            super::checksum(&mut cursor, None::<fn(u64) -> futures::future::Ready<()>>);
+        let mut future = super::checksum(
+            &mut cursor,
+            None::<fn(u64) -> futures::future::Ready<()>>,
+            None,
+        );
         let mut future = unsafe { Pin::new_unchecked(&mut future) };
 
         // expect it to yield 3 times (one at the very end)
diff --git a/drop-transfer/src/ws/client/v4.rs b/drop-transfer/src/ws/client/v4.rs
@@ -140,6 +140,7 @@ impl HandlerLoop<'_> {
                     .checksum::<_, futures::future::Ready<()>>(
                         limit,
                         None::<fn(u64) -> futures::future::Ready<()>>,
+                        None,
                     )
                     .await?;
 
diff --git a/drop-transfer/src/ws/client/v6.rs b/drop-transfer/src/ws/client/v6.rs
@@ -172,6 +172,7 @@ impl HandlerLoop<'_> {
                     .checksum::<_, futures::future::Ready<()>>(
                         limit,
                         None::<fn(u64) -> futures::future::Ready<()>>,
+                        None,
                     )
                     .await?;
 
diff --git a/drop-transfer/src/ws/server/handler.rs b/drop-transfer/src/ws/server/handler.rs
@@ -82,6 +82,7 @@ pub trait Downloader {
         &mut self,
         location: &Hidden<PathBuf>,
         progress_cb: Option<F>,
+        event_granularity: Option<u64>,
     ) -> crate::Result<()>
     where
         F: FnMut(u64) -> Fut + Send + Sync,
diff --git a/drop-transfer/src/ws/server/mod.rs b/drop-transfer/src/ws/server/mod.rs
@@ -712,6 +712,7 @@ impl FileXferTask {
         downloader: &mut impl Downloader,
         offset: u64,
         emit_checksum_events: bool,
+        checksum_events_granularity: u64,
     ) -> crate::Result<PathBuf> {
         let mut out_file = match downloader.open(tmp_loc).await {
             Ok(out_file) => out_file,
@@ -763,21 +764,27 @@ impl FileXferTask {
 
             if emit_checksum_events {
                 events.checksum_start(self.file.size()).await;
-
                 let progress_cb = {
-                    move |progress: u64| async move {
-                        events.checksum_progress(progress).await;
+                    move |progress_bytes: u64| async move {
+                        events.checksum_progress(progress_bytes).await;
                     }
                 };
 
-                downloader.validate(tmp_loc, Some(progress_cb)).await?;
+                downloader
+                    .validate(
+                        tmp_loc,
+                        Some(progress_cb),
+                        Some(checksum_events_granularity),
+                    )
+                    .await?;
 
                 events.checksum_finish().await;
             } else {
                 downloader
                     .validate::<_, futures::future::Ready<()>>(
                         tmp_loc,
                         None::<fn(u64) -> futures::future::Ready<()>>,
+                        None,
                     )
                     .await?;
             }
@@ -865,6 +872,7 @@ impl FileXferTask {
         events: &FileEventTx<IncomingTransfer>,
         tmp_location: &Hidden<PathBuf>,
         emit_checksum_events: bool,
+        checksum_events_granularity: u64,
     ) -> Option<TmpFileState> {
         // TODO: we load the file's metadata to check if we should emit checksum events
         // based on size threshold. However TmpFileState::load also does the
@@ -880,13 +888,19 @@ impl FileXferTask {
             let size = tmp_size.unwrap_or(0);
             events.checksum_start(size).await;
 
-            Some(|progress| events.checksum_progress(progress))
+            Some(|progress_bytes| events.checksum_progress(progress_bytes))
         } else {
             None
         };
 
         // Check if we can resume the temporary file
-        let tmp_file_state = match TmpFileState::load(&tmp_location.0, cb).await {
+        let tmp_file_state = match TmpFileState::load(
+            &tmp_location.0,
+            cb,
+            Some(checksum_events_granularity),
+        )
+        .await
+        {
             Ok(tmp_file_state) => {
                 debug!(
                     logger,
@@ -930,6 +944,7 @@ impl FileXferTask {
                     false
                 }
             };
+            let checksum_events_granularity = state.config.checksum_events_granularity;
 
             events.preflight().await;
 
@@ -939,7 +954,13 @@ impl FileXferTask {
             );
 
             let tmp_file_state = self
-                .handle_tmp_file(&logger, &events, &tmp_location, emit_checksum_events)
+                .handle_tmp_file(
+                    &logger,
+                    &events,
+                    &tmp_location,
+                    emit_checksum_events,
+                    checksum_events_granularity,
+                )
                 .await;
 
             let init_res = downloader.init(&self, tmp_file_state).await?;
@@ -970,6 +991,7 @@ impl FileXferTask {
                         &mut downloader,
                         offset,
                         emit_checksum_events,
+                        checksum_events_granularity,
                     )
                     .await
                 }
@@ -1043,7 +1065,11 @@ impl FileXferTask {
 
 impl TmpFileState {
     // Blocking operation
-    async fn load<F, Fut>(path: &Path, progress_cb: Option<F>) -> io::Result<Self>
+    async fn load<F, Fut>(
+        path: &Path,
+        progress_cb: Option<F>,
+        event_granularity: Option<u64>,
+    ) -> io::Result<Self>
     where
         F: Fn(u64) -> Fut + Sync + Send,
         Fut: Future<Output = ()>,
@@ -1052,7 +1078,7 @@ impl TmpFileState {
 
         let meta = file.metadata()?;
 
-        let csum = file::checksum(file, progress_cb).await?;
+        let csum = file::checksum(file, progress_cb, event_granularity).await?;
         Ok(TmpFileState { meta, csum })
     }
 }
diff --git a/drop-transfer/src/ws/server/v2.rs b/drop-transfer/src/ws/server/v2.rs
@@ -455,7 +455,12 @@ impl handler::Downloader for Downloader {
         .await
     }
 
-    async fn validate<F, Fut>(&mut self, _path: &Hidden<PathBuf>, _: Option<F>) -> crate::Result<()>
+    async fn validate<F, Fut>(
+        &mut self,
+        _path: &Hidden<PathBuf>,
+        _: Option<F>,
+        _: Option<u64>,
+    ) -> crate::Result<()>
     where
         F: FnMut(u64) -> Fut + Send + Sync,
         Fut: Future<Output = ()>,
diff --git a/drop-transfer/src/ws/server/v4.rs b/drop-transfer/src/ws/server/v4.rs
@@ -594,13 +594,14 @@ impl handler::Downloader for Downloader {
         &mut self,
         path: &Hidden<PathBuf>,
         progress_cb: Option<F>,
+        event_granularity: Option<u64>,
     ) -> crate::Result<()>
     where
         F: FnMut(u64) -> Fut + Send + Sync,
         Fut: Future<Output = ()> + Send + Sync,
     {
         let file = std::fs::File::open(&path.0)?;
-        let csum = file::checksum(file, progress_cb).await?;
+        let csum = file::checksum(file, progress_cb, event_granularity).await?;
 
         if self.full_csum.get().await != csum {
             return Err(crate::Error::ChecksumMismatch);
diff --git a/drop-transfer/src/ws/server/v6.rs b/drop-transfer/src/ws/server/v6.rs
@@ -653,13 +653,14 @@ impl handler::Downloader for Downloader {
         &mut self,
         path: &Hidden<PathBuf>,
         progress_cb: Option<F>,
+        event_granularity: Option<u64>,
     ) -> crate::Result<()>
     where
         F: FnMut(u64) -> Fut + Send + Sync,
         Fut: Future<Output = ()> + Send,
     {
         let file = std::fs::File::open(&path.0)?;
-        let csum = file::checksum(file, progress_cb).await?;
+        let csum = file::checksum(file, progress_cb, event_granularity).await?;
 
         if self.full_csum.get().await != csum {
             return Err(crate::Error::ChecksumMismatch);
diff --git a/norddrop/src/device/types.rs b/norddrop/src/device/types.rs
@@ -157,6 +157,12 @@ pub struct Config {
     #[serde(rename = "checksum_events_size_threshold_bytes")]
     pub checksum_events_size_threshold: Option<usize>,
 
+    #[serde(
+        rename = "checksum_events_granularity_bytes",
+        default = "Config::default_checksum_events_granularity"
+    )]
+    pub checksum_events_granularity: u64,
+
     #[serde(default = "Config::default_connection_retries")]
     pub connection_retries: u32,
 }
@@ -165,6 +171,10 @@ impl Config {
     const fn default_connection_retries() -> u32 {
         5
     }
+
+    const fn default_checksum_events_granularity() -> u64 {
+        256 * 1024
+    }
 }
 
 impl From<&drop_transfer::Error> for Status {
@@ -415,6 +425,7 @@ impl From<Config> for drop_config::Config {
             storage_path,
             moose_app_version,
             checksum_events_size_threshold,
+            checksum_events_granularity,
             connection_retries,
         } = val;
 
@@ -425,6 +436,7 @@ impl From<Config> for drop_config::Config {
                 storage_path,
                 checksum_events_size_threshold,
                 connection_retries,
+                checksum_events_granularity,
             },
             moose: drop_config::MooseConfig {
                 app_version: moose_app_version,
@@ -453,7 +465,8 @@ mod tests {
           "storage_path": ":memory:",
           "max_uploads_in_flight": 16,
           "max_requests_per_sec": 15,
-          "checksum_events_size_threshold_bytes": 1234
+          "checksum_events_size_threshold_bytes": 1234,
+          "checksum_events_granularity_bytes": 1024
         }
         "#;
 
@@ -466,6 +479,7 @@ mod tests {
                     transfer_file_limit,
                     storage_path,
                     checksum_events_size_threshold: checksum_events_size_threshold_bytes,
+                    checksum_events_granularity: checksum_events_granularity_bytes,
                     connection_retries,
                 },
             moose:
@@ -483,6 +497,7 @@ mod tests {
         assert_eq!(app_version, "1.2.5");
         assert!(prod);
         assert_eq!(checksum_events_size_threshold_bytes, Some(1234));
+        assert_eq!(checksum_events_granularity_bytes, 1024);
         assert_eq!(connection_retries, 5);
     }
 }
diff --git a/test/drop_test/action.py b/test/drop_test/action.py
@@ -740,16 +740,19 @@ def __init__(
         addr: str,
         dbpath: str = ":memory:",
         checksum_events_size_threshold=2**32,  # don't emit events for existing tests
+        checksum_events_granularity=None,
     ):
         self._addr = addr
         self._dbpath = dbpath
         self._checksum_events_size_threshold = checksum_events_size_threshold
+        self._checksum_events_granularity = checksum_events_granularity
 
     async def run(self, drop: ffi.Drop):
         drop.start(
             peer_resolver.resolve(self._addr),
             self._dbpath,
             self._checksum_events_size_threshold,
+            self._checksum_events_granularity,
         )
 
     def __str__(self):
diff --git a/test/drop_test/ffi.py b/test/drop_test/ffi.py
diff --git a/test/scenarios.py b/test/scenarios.py

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,9 @@ pub struct DropConfig {`
`13`	`13`	`pub storage_path: String,`
`14`	`14`	`// If set the checksum events will be emited for every file of this or bigger size`
`15`	`15`	`pub checksum_events_size_threshold: Option<usize>,`
	`16`	`+ // If set the checksum events will be emited for every checksum_events_granularity bytes`
	`17`	`+ // Default value is 256KB.`
	`18`	`+ pub checksum_events_granularity: u64,`
`16`	`19`	`pub connection_retries: u32,`
`17`	`20`	`}`
`18`	`21`
`@@ -23,6 +26,7 @@ impl Default for DropConfig {`
`23`	`26`	`transfer_file_limit: 1000,`
`24`	`27`	`storage_path: "libdrop.sqlite".to_string(),`
`25`	`28`	`checksum_events_size_threshold: None,`
	`29`	`+ checksum_events_granularity: 256 * 1024,`
`26`	`30`	`connection_retries: 5,`
`27`	`31`	`}`
`28`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -140,6 +140,7 @@ impl HandlerLoop<'_> {`
`140`	`140`	`.checksum::<_, futures::future::Ready<()>>(`
`141`	`141`	`limit,`
`142`	`142`	`None::<fn(u64) -> futures::future::Ready<()>>,`
	`143`	`+ None,`
`143`	`144`	`)`
`144`	`145`	`.await?;`
`145`	`146`
Original file line number	Diff line number	Diff line change
`@@ -172,6 +172,7 @@ impl HandlerLoop<'_> {`
`172`	`172`	`.checksum::<_, futures::future::Ready<()>>(`
`173`	`173`	`limit,`
`174`	`174`	`None::<fn(u64) -> futures::future::Ready<()>>,`
	`175`	`+ None,`
`175`	`176`	`)`
`176`	`177`	`.await?;`
`177`	`178`