fix: Prevent duplicates across magnets/torrents

Author: angrynode

src/database/magnet.rs

@@ -6,6 +6,7 @@ use snafu::prelude::*;
 
 use crate::database::operation::*;
 use crate::database::operator::DatabaseOperator;
+use crate::database::torrent::TorrentError;
 use crate::database::{category, content_folder};
 use crate::extractors::user::User;
 use crate::routes::magnet::MagnetForm;
@@ -53,6 +54,12 @@ pub enum MagnetError {
     NoSuchCategory { id: i32 },
     #[snafu(display("Requested content folder not found"))]
     NoSuchContentFolder { id: i32 },
+    #[snafu(display("The magnet is already uploaded"))]
+    DuplicateMagnet,
+    #[snafu(display("This magnet is already known in the system with the full torrent"))]
+    DuplicateTorrent,
+    #[snafu(display("Failed to read the torrent list"))]
+    Torrent { source: TorrentError },
 }
 
 #[derive(Clone, Debug)]
@@ -200,7 +207,15 @@ impl MagnetOperator {
 
         if list.iter().any(|x| x.torrent_id == magnet.id()) {
             // The magnet is already known
-            return self.get_by_torrent_id(&magnet.id()).await;
+            return Err(MagnetError::DuplicateMagnet);
+        }
+
+        // Check for duplicates in the torrent table, so we don't
+        // even have to resolve the magnet.
+        let list = self.db().torrent().list().await.context(TorrentSnafu)?;
+        if list.iter().any(|x| x.torrent_id == magnet.id()) {
+            // The magnet is already known as a torrent
+            return Err(MagnetError::DuplicateTorrent);
         }
 
         // Verify that the requested category/content_folder exist
@@ -264,4 +279,14 @@ impl MagnetOperator {
 
         Ok(model)
     }
+
+    /// Removes (and cancels resolution) for a given magnet,
+    ///
+    /// if it was previously known. Has no effect otherwise.
+    pub async fn cancel_and_remove(&self, torrent_id: &TorrentID) {
+        if let Ok(magnet) = self.get_by_torrent_id(torrent_id).await {
+            // TODO: should we error here?
+            self.delete(magnet.id).await.unwrap();
+        }
+    }
 }

src/database/torrent.rs

@@ -53,6 +53,8 @@ pub enum TorrentError {
     NoSuchCategory { id: i32 },
     #[snafu(display("Requested content folder not found"))]
     NoSuchContentFolder { id: i32 },
+    #[snafu(display("This torrent is already known"))]
+    DuplicateTorrent,
 }
 
 #[derive(Clone, Debug)]
@@ -203,10 +205,9 @@ impl TorrentOperator {
 
         // Check duplicates
         let list = self.list().await?;
-
         if list.iter().any(|x| x.torrent_id == torrent.id()) {
             // The torrent is already known
-            return self.get_by_torrent_id(&torrent.id()).await;
+            return Err(TorrentError::DuplicateTorrent);
         }
 
         // Verify that the requested category/content_folder exist
@@ -240,6 +241,9 @@ impl TorrentOperator {
         .await
         .context(DBSnafu)?;
 
+        // Now if there was a magnet with this torrentID, we just remove it and cancel the task
+        self.db().magnet().cancel_and_remove(&torrent.id()).await;
+
         // Should not fail
         let model = model.try_into_model().unwrap();
 

src/resolver.rs

@@ -7,7 +7,8 @@ use tokio::time::{Duration, sleep};
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::database::magnet::MagnetOperator;
+use crate::database::magnet;
+use crate::routes::torrent::TorrentForm;
 use crate::state::AppState;
 
 /// A magnet link resolver
@@ -21,11 +22,10 @@ use crate::state::AppState;
 /// for updates on a channel determined on startup. This will
 /// avoid polluting the logs with queries to read the table.
 pub struct Resolver {
-    operator: MagnetOperator,
+    operator: magnet::MagnetOperator,
     // Channel to receive new magnets to resolve
     receiver: UnboundedReceiver<MagnetLink>,
     session: Arc<Session>,
-    state: AppState,
     // Keep track of background tasks resolving torrent files from magnets
     // In the future, this will allow to cancel/delete tasks.
     tasks: HashMap<TorrentID, JoinHandle<TorrentFile>>,
@@ -46,13 +46,12 @@ impl Resolver {
         .unwrap();
 
         Self {
-            operator: MagnetOperator {
+            operator: magnet::MagnetOperator {
                 state: state.clone(),
                 user: None,
             },
             receiver,
             session,
-            state,
             tasks: HashMap::new(),
         }
     }
@@ -131,7 +130,18 @@ impl Resolver {
         }
     }
 
-    /// Check if some tasks have finished resolving, and save result in the DB.
+    /// Check if some magnets have finished resolving, and save result in the DB.
+    ///
+    /// In order to avoid losing data, we check in this order:
+    ///
+    /// - if the magnet no longer exists, do nothing
+    /// - if there is already a corresponding torrent, do nothing
+    /// - save the corresponding torrent
+    /// - remove the corresponding magnet if the saving was successful
+    ///
+    /// On startup, we will further check that no dangling magnets are left due to
+    /// the program crashing/stopping between saving the torrent and removing
+    /// the magnet.
     pub async fn save_resolved(&mut self) {
         let finished_ids: Vec<TorrentID> = self
             .tasks
@@ -147,6 +157,22 @@ impl Resolver {
         for finished_id in finished_ids {
             log::info!("Magnet {finished_id} has finished resolving. Saving to DB.");
 
+            // First verify if someone has in the meantime added the fully resolved torrent?
+            if let Ok(_torrent) = self
+                .operator
+                .db()
+                .torrent()
+                .get_by_torrent_id(&finished_id)
+                .await
+            {
+                // Nothing to do, everything is already well and good.
+                // Well actually, the imported category/folder for the torrent may be different from
+                // the one requested on the magnet. But since we already have a duplicate check on
+                // magnet/torrent upload and this case is only for TOCTOU race cases, ignoring
+                // it seems reasonable.
+                return;
+            }
+
             // Get the raw task handle, removing it from the active tasks
             let handle = self.tasks.remove(&finished_id).unwrap();
             let torrent_file = handle.await.unwrap();
@@ -164,10 +190,22 @@ impl Resolver {
                     continue
                 }
 
-            }
+                // Now that we have saved the torrent state, move to the torrent table
+                let torrent_form = TorrentForm {
+                    category_id: magnet.category_id,
+                    content_folder_id: magnet.content_folder_id,
+                    file: axum::body::Bytes::copy_from_slice(&torrent_file.to_vec()),
+                };
+
+                if let Err(e) = self.operator.db().torrent().create(&torrent_form).await {
+                    // If the creation was not successful, log it!
+                    log::error!(
+                        "Failed to move finished magnet {finished_id} to the torrent table: {e}"
+                    );
+                }
 
-            // TODO: save to the torrent DB
-            log::info!("Torrent file for magnet {finished_id} has been saved to DB");
+                log::info!("Torrent file for magnet {finished_id} has been saved to DB");
+            }
         }
     }
 }