Merge pull request #337 from NordSecurity/msz/FILE-534-implement-auto…

…matic-transfer-cancellation

Implement automatic transfer cancelation

changelog.md

@@ -3,6 +3,7 @@
 ---
 * Remove support for deprecated, unsecure protocols V1, V2, V4 and V5
 * Sanitize file ID before downloading a file, to prevent the temporary file from being allowed to traverse parent dirs
+* Implement automatic transfer cancellation when all files reach the terminal state
 
 ---
 <br>

drop-storage/src/lib.rs

@@ -377,30 +377,19 @@ impl Storage {
         }
     }
 
-    pub async fn update_incoming_file_sync_states(
-        &self,
-        transfer_id: Uuid,
-        file_id: &str,
-        local: sync::FileState,
-    ) {
-        let task = async {
-            let conn = self.conn.lock().await;
-            sync::incoming_file_set_local_state(&conn, transfer_id, file_id, local)?;
-            Ok::<(), Error>(())
-        };
-
-        if let Err(e) = task.await {
+    pub async fn stop_incoming_file(&self, transfer_id: Uuid, file_id: &str) -> Option<()> {
+        let conn = self.conn.lock().await;
+
+        if let Err(e) = sync::incoming_file_set_local_state(
+            &conn,
+            transfer_id,
+            file_id,
+            sync::FileState::Terminal,
+        ) {
             error!(self.logger, "Failed to update incoming file sync states"; "error" => %e);
         }
-    }
 
-    pub async fn stop_incoming_file(&self, transfer_id: Uuid, file_id: &str) -> Option<()> {
-        let task = async {
-            let conn = self.conn.lock().await;
-            sync::stop_incoming_file(&conn, transfer_id, file_id)
-        };
-
-        match task.await {
+        match sync::stop_incoming_file(&conn, transfer_id, file_id) {
             Ok(state) => state,
             Err(e) => {
                 error!(self.logger, "Failed to stop incoming file sync state"; "error" => %e);

drop-transfer/examples/udrop.rs

@@ -1,5 +1,4 @@
 use std::{
-    collections::{btree_map::Entry, BTreeMap, HashSet},
     env,
     io::Write,
     net::IpAddr,
@@ -202,93 +201,19 @@ async fn listen(
 ) -> anyhow::Result<()> {
     info!("Awaiting events…");
 
-    let mut active_file_downloads = BTreeMap::new();
     let mut storage = drop_transfer::StorageDispatch::new(storage);
     while let Some((ev, _)) = rx.recv().await {
         storage.handle_event(&ev).await;
         print_event(&ev);
-        match ev {
-            Event::RequestReceived(xfer) => {
-                let xfid = xfer.id();
-                let files = xfer.files();
-
-                if files.is_empty() {
-                    service
-                        .cancel_all(xfid)
-                        .await
-                        .context("Failed to cancled transfer")?;
-                }
-
-                for file in xfer.files().values() {
-                    service
-                        .download(xfid, file.id(), &out_dir.to_string_lossy())
-                        .await
-                        .context("Cannot issue download call")?;
-                }
-            }
-            Event::FileDownloadStarted(xfer, file, _, _) => {
-                active_file_downloads
-                    .entry(xfer.id())
-                    .or_insert_with(HashSet::new)
-                    .insert(file);
-            }
 
-            Event::FileDownloadProgress(xfer, file, _) => {
-                active_file_downloads
-                    .entry(xfer.id())
-                    .or_insert_with(HashSet::new)
-                    .insert(file);
-            }
-            Event::FileDownloadSuccess(xfer, info) => {
-                let xfid = xfer.id();
-                if let Entry::Occupied(mut occ) = active_file_downloads.entry(xfer.id()) {
-                    occ.get_mut().remove(&info.id);
-                    if occ.get().is_empty() {
-                        service
-                            .cancel_all(xfid)
-                            .await
-                            .context("Failed to cancled transfer")?;
-                        occ.remove_entry();
-                    }
-                }
-            }
-            Event::FileDownloadFailed(xfer, file, _) => {
-                let xfid = xfer.id();
-
-                if let Entry::Occupied(mut occ) = active_file_downloads.entry(xfer.id()) {
-                    occ.get_mut().remove(&file);
-                    if occ.get().is_empty() {
-                        service
-                            .cancel_all(xfid)
-                            .await
-                            .context("Failed to cancled transfer")?;
-                        occ.remove_entry();
-                    }
-                }
-            }
-            Event::FileDownloadRejected {
-                transfer_id,
-                file_id,
-                ..
-            } => {
-                if let Entry::Occupied(mut occ) = active_file_downloads.entry(transfer_id) {
-                    occ.get_mut().remove(&file_id);
-                    if occ.get().is_empty() {
-                        service
-                            .cancel_all(transfer_id)
-                            .await
-                            .context("Failed to cancled transfer")?;
-                        occ.remove_entry();
-                    }
-                }
-            }
-            Event::IncomingTransferCanceled(xfer, _) => {
-                active_file_downloads.remove(&xfer.id());
-            }
-            Event::OutgoingTransferCanceled(xfer, _) => {
-                active_file_downloads.remove(&xfer.id());
+        if let Event::RequestReceived(xfer) = ev {
+            let xfid = xfer.id();
+            for file in xfer.files().values() {
+                service
+                    .download(xfid, file.id(), &out_dir.to_string_lossy())
+                    .await
+                    .context("Cannot issue download call")?;
             }
-            _ => (),
         }
     }
 

drop-transfer/src/check.rs

@@ -34,20 +34,11 @@ pub(crate) fn spawn(
             utils::RetryTrigger::new(refresh_trigger, state.config.connection_retries);
 
         let task = async {
-            loop {
-                let cf = run(&state, &xfer, &logger).await;
-                if cf.is_break() {
-                    info!(logger, "Transfer {} is gone. Clearing", xfer.id());
-
-                    if let Some(state) = state.transfer_manager.incoming_remove(xfer.id()).await {
-                        state.xfer_events.cancel(true).await
-                    }
-
-                    break;
-                }
-
+            while run(&state, &xfer, &logger).await.is_continue() {
                 backoff.backoff().await;
             }
+
+            info!(logger, "Transfer {} is gone. Clearing", xfer.id());
         };
 
         tokio::select! {
@@ -68,11 +59,18 @@ async fn run(state: &State, xfer: &Arc<IncomingTransfer>, logger: &Logger) -> Co
         return ControlFlow::Break(());
     }
 
-    ask_server(state, xfer, logger).await?;
+    if !ask_server_if_alive(state, xfer, logger).await {
+        if let Some(state) = state.transfer_manager.incoming_remove(xfer.id()).await {
+            state.xfer_events.cancel(true).await
+        }
+
+        return ControlFlow::Break(());
+    }
+
     ControlFlow::Continue(())
 }
 
-async fn ask_server(state: &State, xfer: &IncomingTransfer, logger: &Logger) -> ControlFlow<()> {
+async fn ask_server_if_alive(state: &State, xfer: &IncomingTransfer, logger: &Logger) -> bool {
     let mut connector = hyper::client::HttpConnector::new();
     connector.set_local_address(Some(state.addr));
 
@@ -91,7 +89,7 @@ async fn ask_server(state: &State, xfer: &IncomingTransfer, logger: &Logger) ->
         )
         .await
         {
-            Ok(false) => return ControlFlow::Break(()),
+            Ok(false) => return false,
             Ok(true) => break,
             Err(RequestError::UnexpectedResponse(status)) => {
                 debug!(
@@ -110,7 +108,7 @@ async fn ask_server(state: &State, xfer: &IncomingTransfer, logger: &Logger) ->
         }
     }
 
-    ControlFlow::Continue(())
+    true
 }
 
 // Returns whether the transfer is alive