refactor: de-indent transaction fns in chat.rs

The issue has been introduced in
cf33db3dcb
(https://github.com/chatmail/core/pull/2887).

This, again, has to do with a race where the draft message
is sent in another Future after `get_draft` but before `sql.execute`.

Related:
- 07fa9c35ee
  (https://github.com/chatmail/core/pull/6052).
- df4fd82140
  (https://github.com/chatmail/core/pull/6061).

src/chat.rs

@@ -748,6 +748,18 @@ SELECT id, rfc724_mid, pre_rfc724_mid, timestamp, ?, 1 FROM msgs WHERE chat_id=?
         Ok(msg_id)
     }
 
+    fn has_draft_with_id(&self, conn: &rusqlite::Connection, draft_id: &MsgId) -> Result<bool> {
+        let count: u32 = conn.query_row(
+            "SELECT
+                COUNT(*)
+            FROM msgs
+            WHERE id=? AND chat_id=? AND state=?",
+            (draft_id, self, MessageState::OutDraft),
+            |row| row.get(0),
+        )?;
+        Ok(count > 0)
+    }
+
     /// Returns draft message, if there is one.
     pub async fn get_draft(self, context: &Context) -> Result<Option<Message>> {
         if self.is_special() {
@@ -819,48 +831,41 @@ SELECT id, rfc724_mid, pre_rfc724_mid, timestamp, ?, 1 FROM msgs WHERE chat_id=?
         msg.state = MessageState::OutDraft;
         msg.chat_id = self;
 
-        // if possible, replace existing draft and keep id
-        if !msg.id.is_special()
-            && let Some(old_draft) = self.get_draft(context).await?
-            && old_draft.id == msg.id
-            && old_draft.chat_id == self
-            && old_draft.state == MessageState::OutDraft
-        {
-            let affected_rows = context
-                        .sql.execute(
-                                "UPDATE msgs
-                                SET timestamp=?1,type=?2,txt=?3,txt_normalized=?4,param=?5,mime_in_reply_to=?6
-                                WHERE id=?7
-                                AND (type <> ?2 
-                                    OR txt <> ?3 
-                                    OR txt_normalized <> ?4
-                                    OR param <> ?5
-                                    OR mime_in_reply_to <> ?6);",
-                                (
-                                    time(),
-                                    msg.viewtype,
-                                    &msg.text,
-                                    normalize_text(&msg.text),
-                                    msg.param.to_string(),
-                                    msg.in_reply_to.as_deref().unwrap_or_default(),
-                                    msg.id,
-                                ),
-                            ).await?;
-            return Ok(affected_rows > 0);
-        }
-
-        let row_id = context
-            .sql
-            .transaction(|transaction| {
-                // Delete existing draft if it exists.
-                transaction.execute(
-                    "DELETE FROM msgs WHERE chat_id=? AND state=?",
-                    (self, MessageState::OutDraft),
+        let trans_fn = |transaction: &mut rusqlite::Transaction| {
+            // if possible, replace existing draft and keep id
+            if !msg.id.is_special() && self.has_draft_with_id(transaction, &msg.id)? {
+                let affected_rows = transaction.execute(
+                    "UPDATE msgs
+                    SET timestamp=?1,type=?2,txt=?3,txt_normalized=?4,param=?5,mime_in_reply_to=?6
+                    WHERE id=?7
+                    AND (type <> ?2 
+                        OR txt <> ?3 
+                        OR txt_normalized <> ?4
+                        OR param <> ?5
+                        OR mime_in_reply_to <> ?6);",
+                    (
+                        time(),
+                        msg.viewtype,
+                        &msg.text,
+                        normalize_text(&msg.text),
+                        msg.param.to_string(),
+                        msg.in_reply_to.as_deref().unwrap_or_default(),
+                        msg.id,
+                    ),
                 )?;
+                let changed = affected_rows > 0;
+                return Ok((msg.id, changed));
+            }
 
-                // Insert new draft.
-                transaction.execute(
-                    "INSERT INTO msgs (
+            // Delete existing draft if it exists.
+            transaction.execute(
+                "DELETE FROM msgs WHERE chat_id=? AND state=?",
+                (self, MessageState::OutDraft),
+            )?;
+
+            // Insert new draft.
+            transaction.execute(
+                "INSERT INTO msgs (
                  chat_id,
                  rfc724_mid,
                  from_id,
@@ -873,26 +878,28 @@ SELECT id, rfc724_mid, pre_rfc724_mid, timestamp, ?, 1 FROM msgs WHERE chat_id=?
                  hidden,
                  mime_in_reply_to)
          VALUES (?,?,?,?,?,?,?,?,?,?,?);",
-                    (
-                        self,
-                        &msg.rfc724_mid,
-                        ContactId::SELF,
-                        time(),
-                        msg.viewtype,
-                        MessageState::OutDraft,
-                        &msg.text,
-                        normalize_text(&msg.text),
-                        msg.param.to_string(),
-                        1,
-                        msg.in_reply_to.as_deref().unwrap_or_default(),
-                    ),
-                )?;
+                (
+                    self,
+                    &msg.rfc724_mid,
+                    ContactId::SELF,
+                    time(),
+                    msg.viewtype,
+                    MessageState::OutDraft,
+                    &msg.text,
+                    normalize_text(&msg.text),
+                    msg.param.to_string(),
+                    1,
+                    msg.in_reply_to.as_deref().unwrap_or_default(),
+                ),
+            )?;
 
-                Ok(transaction.last_insert_rowid())
-            })
-            .await?;
-        msg.id = MsgId::new(row_id.try_into()?);
-        Ok(true)
+            let msg_id = MsgId::new(transaction.last_insert_rowid().try_into()?);
+            let changed = true;
+            Ok((msg_id, changed))
+        };
+        let (msg_id, changed) = context.sql.transaction(trans_fn).await?;
+        msg.id = msg_id;
+        Ok(changed)
     }
 
     /// Returns number of messages in a chat.
@@ -1758,13 +1765,16 @@ impl Chat {
     /// Adds missing values to the msg object,
     /// writes the record to the database.
     ///
-    /// If `update_msg_id` is set, that record is reused;
-    /// if `update_msg_id` is None, a new record is created.
+    /// If `update_existing_draft == `[`UseExistingDraftPolicy::Reuse`],
+    /// the existing draft with ID == msg.id is reused.
+    /// If no such draft exists, an error is returned.
+    /// If `update_existing_draft == `[`UseExistingDraftPolicy::DontReuse`],
+    /// a new record is created.
     async fn prepare_msg_raw(
         &mut self,
         context: &Context,
         msg: &mut Message,
-        update_msg_id: Option<MsgId>,
+        update_existing_draft: UseExistingDraftPolicy,
     ) -> Result<()> {
         let mut to_id = 0;
         let mut location_id = 0;
@@ -1933,10 +1943,24 @@ impl Chat {
         msg.from_id = ContactId::SELF;
 
         // add message to the database
-        if let Some(update_msg_id) = update_msg_id {
-            context
-                .sql
-                .execute(
+        let trans_fn = |transaction: &mut rusqlite::Transaction| {
+            if update_existing_draft == UseExistingDraftPolicy::Reuse {
+                // This check also covers the `msg.id.is_special()` case.
+                // Maybe we could try to somehow gracefully recover from this,
+                // but better safe than sorry.
+                if !self.id.has_draft_with_id(transaction, &msg.id)? {
+                    bail!(
+                        concat!(
+                            "wanted to prepare existing draft for sending in chat {0}, ",
+                            "but no draft with ID {1} is present ",
+                            "(it might have been sent or deleted)"
+                        ),
+                        self.id,
+                        msg.id
+                    );
+                }
+
+                transaction.execute(
                     "UPDATE msgs
                      SET rfc724_mid=?, chat_id=?, from_id=?, to_id=?, timestamp=?, type=?,
                          state=?, txt=?, txt_normalized=?, subject=?, param=?,
@@ -1964,37 +1988,35 @@ impl Chat {
                         location_id as i32,
                         ephemeral_timer,
                         ephemeral_timestamp,
-                        update_msg_id
+                        msg.id
                     ],
-                )
-                .await?;
-            msg.id = update_msg_id;
-        } else {
-            let raw_id = context
-                .sql
-                .insert(
+                )?;
+                let inserted = false;
+                Ok((msg.id, inserted))
+            } else {
+                transaction.execute(
                     "INSERT INTO msgs (
-                        rfc724_mid,
-                        chat_id,
-                        from_id,
-                        to_id,
-                        timestamp,
-                        type,
-                        state,
-                        txt,
-                        txt_normalized,
-                        subject,
-                        param,
-                        hidden,
-                        mime_in_reply_to,
-                        mime_references,
-                        mime_modified,
-                        mime_headers,
-                        mime_compressed,
-                        location_id,
-                        ephemeral_timer,
-                        ephemeral_timestamp)
-                        VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,1,?,?,?);",
+                     rfc724_mid,
+                     chat_id,
+                     from_id,
+                     to_id,
+                     timestamp,
+                     type,
+                     state,
+                     txt,
+                     txt_normalized,
+                     subject,
+                     param,
+                     hidden,
+                     mime_in_reply_to,
+                     mime_references,
+                     mime_modified,
+                     mime_headers,
+                     mime_compressed,
+                     location_id,
+                     ephemeral_timer,
+                     ephemeral_timestamp)
+                     VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,1,?,?,?);",
                     params_slice![
                         msg.rfc724_mid,
                         msg.chat_id,
@@ -2016,10 +2038,17 @@ impl Chat {
                         ephemeral_timer,
                         ephemeral_timestamp
                     ],
-                )
-                .await?;
+                )?;
+                let msg_id = MsgId::new(transaction.last_insert_rowid().try_into()?);
+                let inserted = true;
+                Ok((msg_id, inserted))
+            }
+        };
+        let (msg_id, inserted) = context.sql.transaction(trans_fn).await?;
+
+        msg.id = msg_id;
+        if inserted {
             context.new_msgs_notify.notify_one();
-            msg.id = MsgId::new(u32::try_from(raw_id)?);
 
             maybe_set_logging_xdc(context, msg, self.id).await?;
             context
@@ -2716,15 +2745,15 @@ async fn prepare_send_msg(
     }
 
     // check current MessageState for drafts (to keep msg_id) ...
-    let update_msg_id = if msg.state == MessageState::OutDraft {
+    let update_existing_draft = if msg.state == MessageState::OutDraft {
         msg.hidden = false;
         if !msg.id.is_special() && msg.chat_id == chat_id {
-            Some(msg.id)
+            UseExistingDraftPolicy::Reuse
         } else {
-            None
+            UseExistingDraftPolicy::DontReuse
         }
     } else {
-        None
+        UseExistingDraftPolicy::DontReuse
     };
 
     if msg.state == MessageState::Undefined
@@ -2744,7 +2773,8 @@ async fn prepare_send_msg(
     if !msg.hidden {
         chat_id.unarchive_if_not_muted(context, msg.state).await?;
     }
-    chat.prepare_msg_raw(context, msg, update_msg_id).await?;
+    chat.prepare_msg_raw(context, msg, update_existing_draft)
+        .await?;
 
     let row_ids = create_send_msg_jobs(context, msg)
         .await
@@ -2755,6 +2785,36 @@ async fn prepare_send_msg(
     Ok(row_ids)
 }
 
+/// When sending a message or setting the draft,
+/// what to do about the potentially existing draft.
+///
+/// This is basically a more explicit bool.
+#[derive(Debug, PartialEq)]
+enum UseExistingDraftPolicy {
+    /// Create a brand new draft or message, don't reuse the existing one's ID.
+    DontReuse,
+    /// Resuse the existing draft, so that the new draft or message
+    /// keeps its original ID.
+    /// Useful for WebXDC attachments.
+    Reuse,
+}
+impl From<UseExistingDraftPolicy> for bool {
+    fn from(reuse: UseExistingDraftPolicy) -> bool {
+        match reuse {
+            UseExistingDraftPolicy::DontReuse => false,
+            UseExistingDraftPolicy::Reuse => true,
+        }
+    }
+}
+impl From<bool> for UseExistingDraftPolicy {
+    fn from(reuse: bool) -> UseExistingDraftPolicy {
+        match reuse {
+            false => UseExistingDraftPolicy::DontReuse,
+            true => UseExistingDraftPolicy::Reuse,
+        }
+    }
+}
+
 /// Renders the Message or splits it into Pre- and Post-Message.
 ///
 /// Pre-Message is a small message with metadata which announces a larger Post-Message.
@@ -4576,7 +4636,8 @@ pub async fn forward_msgs_2ctx(
         msg.rfc724_mid = create_outgoing_rfc724_mid();
         msg.pre_rfc724_mid.clear();
         msg.timestamp_sort = now;
-        chat.prepare_msg_raw(ctx_dst, &mut msg, None).await?;
+        chat.prepare_msg_raw(ctx_dst, &mut msg, UseExistingDraftPolicy::DontReuse)
+            .await?;
 
         if !create_send_msg_jobs(ctx_dst, &mut msg).await?.is_empty() {
             ctx_dst.scheduler.interrupt_smtp().await;

src/chat/chat_tests.rs

@@ -175,6 +175,41 @@ async fn test_draft_stable_ids() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_dont_send_sent_draft() -> Result<()> {
+    let mut tcm = TestContextManager::new();
+    let t = &tcm.alice().await;
+    let bob = &tcm.bob().await;
+    let chat_id = t.create_chat(bob).await.id;
+
+    let mut msg = Message::new_text("original".to_string());
+
+    chat_id.set_draft(t, Some(&mut msg)).await?;
+    assert_eq!(msg.state, MessageState::OutDraft);
+
+    let mut msg_clone = msg.clone();
+
+    send_msg(t, chat_id, &mut msg).await?;
+    assert_eq!(msg.state, MessageState::OutPending);
+
+    msg_clone.set_text("modified".to_string());
+    // Try to send the stale draft Message object with the same ID again.
+    assert_eq!(msg_clone.id, msg.id);
+    assert_eq!(msg_clone.state, MessageState::OutDraft);
+    let msg_from_db_before_send = Message::load_from_db(t, msg.id).await?;
+    let send_res = send_msg(t, chat_id, &mut msg_clone).await;
+    assert!(send_res.is_err());
+
+    let msg_from_db_after_send = Message::load_from_db(t, msg.id).await?;
+    assert_eq!(msg_from_db_after_send.text, "original");
+    assert_eq!(
+        serde_json::to_string_pretty(&msg_from_db_before_send).unwrap(),
+        serde_json::to_string_pretty(&msg_from_db_after_send).unwrap()
+    );
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_only_one_draft_per_chat() -> Result<()> {
     let t = TestContext::new_alice().await;