From 9efcde6f57dbf9b4fec139b258713ad88603e969 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Fri, 16 Oct 2020 22:29:21 +0100
Subject: [PATCH] fix: killing subscriber when smp connection is SUBd in
 another tcp connection

---
 src/Server.hs | 17 ++++++++---------
 tests/Test.hs | 19 +++++++++++++++++--
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/Server.hs b/src/Server.hs
index a802f1c38..7450f3fb7 100644
--- a/src/Server.hs
+++ b/src/Server.hs
@@ -26,7 +26,6 @@ import Transmission
 import Transport
 import UnliftIO.Async
 import UnliftIO.Concurrent
-import qualified UnliftIO.Exception as E
 import UnliftIO.IO
 import UnliftIO.STM
 
@@ -168,7 +167,7 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
           ms <- asks msgStore
           q <- getMsgQueue ms rId
           tryPeek q >>= \case
-            Just msg -> return $ msgResponse msg
+            Just msg -> return $ msgResponse rId msg
             Nothing -> forkSubscriber q rId
 
         forkSubscriber :: MsgQueue -> RecipientId -> m Signed
@@ -176,15 +175,15 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
           cs <- readTVarIO connections
           case M.lookup rId cs of
             Just (Left ()) -> do
-              E.bracket
-                (forkIO subscriber)
-                (\_ -> trackSubscriber $ Left ())
-                (trackSubscriber . Right)
+              threadId <- forkIO subscriber
+              trackSubscriber $ Right threadId
               return ok
             _ -> return ok
           where
             trackSubscriber sThrd = atomically . modifyTVar connections $ M.insert rId sThrd
-            subscriber = peekMsg q >>= atomically . writeTBQueue sndQ . msgResponse
+            subscriber = do
+              peekMsg q >>= atomically . writeTBQueue sndQ . msgResponse rId
+              trackSubscriber $ Left ()
 
         ok :: Signed
         ok = (connId, Cmd SBroker OK)
@@ -195,5 +194,5 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
         okResponse :: Either ErrorType () -> Signed
         okResponse = mkSigned connId . either ERR (const OK)
 
-        msgResponse :: Message -> Signed
-        msgResponse Message {msgId, ts, msgBody} = mkSigned connId $ MSG msgId ts msgBody
+        msgResponse :: RecipientId -> Message -> Signed
+        msgResponse rId Message {msgId, ts, msgBody} = mkSigned rId $ MSG msgId ts msgBody
diff --git a/tests/Test.hs b/tests/Test.hs
index 5aa352007..7abb2f41b 100644
--- a/tests/Test.hs
+++ b/tests/Test.hs
@@ -44,6 +44,12 @@ testCreateSecure = do
       (ok1, OK) #== "accepts unsigned SEND"
       (sId1, sId) #== "same connection ID in response 1"
 
+      Resp _ (MSG _ _ msg1) <- tGet fromServer h
+      (msg1, "hello") #== "delivers message"
+
+      Resp _ ok4 <- sendRecv h ("123", "1", "ACK")
+      (ok4, OK) #== "replies OK when message acknowledged if no more messages"
+
       Resp sId2 err1 <- sendRecv h ("456", sId, "SEND :hello")
       (err1, ERR AUTH) #== "rejects signed SEND"
       (sId2, sId) #== "same connection ID in response 2"
@@ -61,9 +67,15 @@ testCreateSecure = do
       Resp _ err4 <- sendRecv h ("123", rId, "KEY 456")
       (err4, ERR AUTH) #== "rejects KEY if already secured"
 
-      Resp _ ok3 <- sendRecv h ("456", sId, "SEND :hello")
+      Resp _ ok3 <- sendRecv h ("456", sId, "SEND 11\nhello again")
       (ok3, OK) #== "accepts signed SEND"
 
+      Resp _ (MSG _ _ msg) <- tGet fromServer h
+      (msg, "hello again") #== "delivers message 2"
+
+      Resp _ ok5 <- sendRecv h ("123", "1", "ACK")
+      (ok5, OK) #== "replies OK when message acknowledged 2"
+
       Resp _ err5 <- sendRecv h ("", sId, "SEND :hello")
       (err5, ERR AUTH) #== "rejects unsigned SEND"
 
@@ -80,6 +92,9 @@ testCreateDelete = do
       Resp _ ok2 <- sendRecv h ("456", sId, "SEND :hello")
       (ok2, OK) #== "accepts signed SEND"
 
+      Resp _ (MSG _ _ msg1) <- tGet fromServer h
+      (msg1, "hello") #== "delivers message"
+
       Resp _ err1 <- sendRecv h ("1234", rId, "OFF")
       (err1, ERR AUTH) #== "rejects OFF with wrong signature (password atm)"
 
@@ -100,7 +115,7 @@ testCreateDelete = do
       (ok4, OK) #== "accepts OFF when suspended"
 
       Resp _ (MSG _ _ msg) <- sendRecv h ("123", rId, "SUB")
-      (msg, "hello") #== "accepts SUB when suspended and delivers the message"
+      (msg, "hello") #== "accepts SUB when suspended and delivers the message again (because was not ACKed)"
 
       Resp _ err5 <- sendRecv h ("1234", rId, "DEL")
       (err5, ERR AUTH) #== "rejects DEL with wrong signature (password atm)"