From 9893935e7c3cf8d102c85730a4e48d32f05c2ec7 Mon Sep 17 00:00:00 2001
From: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>
Date: Thu, 5 Dec 2024 19:27:21 +0400
Subject: [PATCH] agent: fix acceptContact inconsistent state after app restart
 (#1412)

---
 src/Simplex/Messaging/Agent.hs              | 17 +++++++++++------
 src/Simplex/Messaging/Agent/Client.hs       |  2 +-
 src/Simplex/Messaging/Agent/Store.hs        |  2 +-
 src/Simplex/Messaging/Agent/Store/SQLite.hs |  6 +++---
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 94db36e28..f267baf28 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -687,9 +687,15 @@ allowConnectionAsync' c corrId connId confId ownConnInfo =
       enqueueCommand c corrId connId (Just server) $ AClientCommand $ LET confId ownConnInfo
     _ -> throwE $ CMD PROHIBITED "allowConnectionAsync"
 
+-- TODO
+-- Unlike `acceptContact` (synchronous version), `acceptContactAsync` uses `unacceptInvitation` in case of error,
+-- because we're not taking lock here. In practice it is less likely to fail because it doesn't involve network IO,
+-- and also it can't be triggered by user concurrently several times in a row. It could be improved similarly to
+-- `acceptContact` by creating a new map for invitation locks and taking lock here, and removing `unacceptInvitation`
+-- while marking invitation as accepted inside "lock level transaction" after successful `joinConnAsync`.
 acceptContactAsync' :: AgentClient -> ACorrId -> Bool -> InvitationId -> ConnInfo -> PQSupport -> SubscriptionMode -> AM ConnId
 acceptContactAsync' c corrId enableNtfs invId ownConnInfo pqSupport subMode = do
-  Invitation {contactConnId, connReq} <- withStore c (`getInvitation` invId)
+  Invitation {contactConnId, connReq} <- withStore c $ \db -> getInvitation db "acceptContactAsync'" invId
   withStore c (`getConn` contactConnId) >>= \case
     SomeConn _ (ContactConnection ConnData {userId} _) -> do
       withStore' c $ \db -> acceptInvitation db invId ownConnInfo
@@ -809,7 +815,7 @@ newConnToJoin c userId connId enableNtfs cReq pqSup = case cReq of
 
 newConnToAccept :: AgentClient -> ConnId -> Bool -> ConfirmationId -> PQSupport -> AM ConnId
 newConnToAccept c connId enableNtfs invId pqSup = do
-  Invitation {connReq, contactConnId} <- withStore c (`getInvitation` invId)
+  Invitation {connReq, contactConnId} <- withStore c $ \db -> getInvitation db "newConnToAccept" invId
   withStore c (`getConn` contactConnId) >>= \case
     SomeConn _ (ContactConnection ConnData {userId} _) ->
       newConnToJoin c userId connId enableNtfs connReq pqSup
@@ -941,13 +947,12 @@ allowConnection' c connId confId ownConnInfo = withConnLock c connId "allowConne
 -- | Accept contact (ACPT command) in Reader monad
 acceptContact' :: AgentClient -> ConnId -> Bool -> InvitationId -> ConnInfo -> PQSupport -> SubscriptionMode -> AM SndQueueSecured
 acceptContact' c connId enableNtfs invId ownConnInfo pqSupport subMode = withConnLock c connId "acceptContact" $ do
-  Invitation {contactConnId, connReq} <- withStore c (`getInvitation` invId)
+  Invitation {contactConnId, connReq} <- withStore c $ \db -> getInvitation db "acceptContact'" invId
   withStore c (`getConn` contactConnId) >>= \case
     SomeConn _ (ContactConnection ConnData {userId} _) -> do
+      sqSecured <- joinConn c userId connId enableNtfs connReq ownConnInfo pqSupport subMode
       withStore' c $ \db -> acceptInvitation db invId ownConnInfo
-      joinConn c userId connId enableNtfs connReq ownConnInfo pqSupport subMode `catchAgentError` \err -> do
-        withStore' c (`unacceptInvitation` invId)
-        throwE err
+      pure sqSecured
     _ -> throwE $ CMD PROHIBITED "acceptContact"
 
 -- | Reject contact (RJCT command) in Reader monad
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index af241c2f4..ef55870fe 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -2020,7 +2020,7 @@ storeError = \case
   SEConnDuplicate -> CONN DUPLICATE
   SEBadConnType CRcv -> CONN SIMPLEX
   SEBadConnType CSnd -> CONN SIMPLEX
-  SEInvitationNotFound -> CMD PROHIBITED "SEInvitationNotFound"
+  SEInvitationNotFound cxt invId -> CMD PROHIBITED $ "SEInvitationNotFound " <> cxt <> ", invitationId = " <> show invId
   -- this error is never reported as store error,
   -- it is used to wrap agent operations when "transaction-like" store access is needed
   -- NOTE: network IO should NOT be used inside AgentStoreMonad
diff --git a/src/Simplex/Messaging/Agent/Store.hs b/src/Simplex/Messaging/Agent/Store.hs
index ae010a884..0e2a7bbe9 100644
--- a/src/Simplex/Messaging/Agent/Store.hs
+++ b/src/Simplex/Messaging/Agent/Store.hs
@@ -623,7 +623,7 @@ data StoreError
   | -- | Confirmation not found.
     SEConfirmationNotFound
   | -- | Invitation not found
-    SEInvitationNotFound
+    SEInvitationNotFound String InvitationId
   | -- | Message not found
     SEMsgNotFound
   | -- | Command not found
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs
index 259678f8f..7e60e4070 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs
@@ -926,9 +926,9 @@ createInvitation db gVar NewInvitation {contactConnId, connReq, recipientConnInf
       |]
       (invitationId, contactConnId, connReq, recipientConnInfo)
 
-getInvitation :: DB.Connection -> InvitationId -> IO (Either StoreError Invitation)
-getInvitation db invitationId =
-  firstRow invitation SEInvitationNotFound $
+getInvitation :: DB.Connection -> String -> InvitationId -> IO (Either StoreError Invitation)
+getInvitation db cxt invitationId =
+  firstRow invitation (SEInvitationNotFound cxt invitationId) $
     DB.query
       db
       [sql|