Merge pull request #2 from simplex-chat/correlation-id

add corellationId to SMP protocol server (WIP)

src/Simplex/Messaging/Server.hs

@@ -51,7 +51,7 @@ runSMPServer cfg@ServerConfig {tcpPort} = do
       (rId, clnt) <- readTBQueue subscribedQ
       cs <- readTVar subscribers
       case M.lookup rId cs of
-        Just Client {rcvQ} -> writeTBQueue rcvQ (rId, Cmd SBroker END)
+        Just Client {rcvQ} -> writeTBQueue rcvQ (B.empty, rId, Cmd SBroker END)
         Nothing -> return ()
       writeTVar subscribers $ M.insert rId clnt cs
 
@@ -75,10 +75,10 @@ cancelSub = \case
 
 receive :: (MonadUnliftIO m, MonadReader Env m) => Handle -> Client -> m ()
 receive h Client {rcvQ} = forever $ do
-  (signature, (queueId, cmdOrError)) <- tGet fromClient h
+  (signature, (corrId, queueId, cmdOrError)) <- tGet fromClient h
   signed <- case cmdOrError of
-    Left e -> return . mkResp queueId $ ERR e
-    Right cmd -> verifyTransmission signature queueId cmd
+    Left e -> return . mkResp corrId queueId $ ERR e
+    Right cmd -> verifyTransmission (signature, (corrId, queueId, cmd))
   atomically $ writeTBQueue rcvQ signed
 
 send :: MonadUnliftIO m => Handle -> Client -> m ()
@@ -86,12 +86,12 @@ send h Client {sndQ} = forever $ do
   signed <- atomically $ readTBQueue sndQ
   tPut h (B.empty, signed)
 
-mkResp :: QueueId -> Command 'Broker -> Signed
-mkResp queueId command = (queueId, Cmd SBroker command)
+mkResp :: CorrelationId -> QueueId -> Command 'Broker -> Signed
+mkResp corrId queueId command = (corrId, queueId, Cmd SBroker command)
 
-verifyTransmission :: forall m. (MonadUnliftIO m, MonadReader Env m) => Signature -> QueueId -> Cmd -> m Signed
-verifyTransmission signature queueId cmd = do
-  (queueId,) <$> case cmd of
+verifyTransmission :: forall m. (MonadUnliftIO m, MonadReader Env m) => Transmission -> m Signed
+verifyTransmission (signature, (corrId, queueId, cmd)) = do
+  (corrId,queueId,) <$> case cmd of
     Cmd SBroker _ -> return $ smpErr INTERNAL -- it can only be client command, because `fromClient` was used
     Cmd SRecipient (NEW _) -> return cmd
     Cmd SRecipient _ -> withQueueRec SRecipient $ verifySignature . recipientKey
@@ -121,11 +121,11 @@ client clnt@Client {subscriptions, rcvQ, sndQ} Server {subscribedQ} =
       >>= atomically . writeTBQueue sndQ
   where
     processCommand :: Signed -> m Signed
-    processCommand (queueId, cmd) = do
+    processCommand (corrId, queueId, cmd) = do
       st <- asks queueStore
       case cmd of
-        Cmd SBroker END -> unsubscribeQueue $> (queueId, cmd)
-        Cmd SBroker _ -> return (queueId, cmd)
+        Cmd SBroker END -> unsubscribeQueue $> (corrId, queueId, cmd)
+        Cmd SBroker _ -> return (corrId, queueId, cmd)
         Cmd SSender (SEND msgBody) -> sendMessage st msgBody
         Cmd SRecipient command -> case command of
           NEW rKey -> createQueue st rKey
@@ -136,7 +136,7 @@ client clnt@Client {subscriptions, rcvQ, sndQ} Server {subscribedQ} =
           DEL -> delQueueAndMsgs st
       where
         createQueue :: QueueStore -> RecipientKey -> m Signed
-        createQueue st rKey = mkResp B.empty <$> addSubscribe
+        createQueue st rKey = mkResp corrId B.empty <$> addSubscribe
           where
             addSubscribe =
               addQueueRetry 3 >>= \case
@@ -217,7 +217,7 @@ client clnt@Client {subscriptions, rcvQ, sndQ} Server {subscribedQ} =
             q <- atomically $ getMsgQueue ms rId
             atomically (tryPeek q) >>= \case
               Nothing -> forkSub q $> ok
-              Just msg -> atomically setDelivered $> msgResp rId msg
+              Just msg -> atomically setDelivered $> mkResp corrId rId (msgCmd msg)
           _ -> return ok
           where
             forkSub :: MsgQueue -> m ()
@@ -231,7 +231,7 @@ client clnt@Client {subscriptions, rcvQ, sndQ} Server {subscribedQ} =
             subscriber :: MsgQueue -> m ()
             subscriber q = atomically $ do
               msg <- peekMsg q
-              writeTBQueue sndQ $ msgResp rId msg
+              writeTBQueue sndQ $ mkResp B.empty rId (msgCmd msg)
               setSub (\s -> s {subThread = NoSub})
               void setDelivered
 
@@ -250,16 +250,16 @@ client clnt@Client {subscriptions, rcvQ, sndQ} Server {subscribedQ} =
               Right _ -> delMsgQueue ms queueId $> ok
 
         ok :: Signed
-        ok = mkResp queueId OK
+        ok = mkResp corrId queueId OK
 
         err :: ErrorType -> Signed
-        err = mkResp queueId . ERR
+        err = mkResp corrId queueId . ERR
 
         okResp :: Either ErrorType () -> Signed
         okResp = either err $ const ok
 
-        msgResp :: RecipientId -> Message -> Signed
-        msgResp rId Message {msgId, ts, msgBody} = mkResp rId $ MSG msgId ts msgBody
+        msgCmd :: Message -> Command 'Broker
+        msgCmd Message {msgId, ts, msgBody} = MSG msgId ts msgBody
 
 randomId :: (MonadUnliftIO m, MonadReader Env m) => Int -> m Encoded
 randomId n = do

src/Simplex/Messaging/Server/Transmission.hs

@@ -33,15 +33,15 @@ data Cmd where
 
 deriving instance Show Cmd
 
-type Signed = (QueueId, Cmd)
+type Signed = (CorrelationId, QueueId, Cmd)
 
 type Transmission = (Signature, Signed)
 
-type SignedOrError = (QueueId, Either ErrorType Cmd)
+type SignedOrError = (CorrelationId, QueueId, Either ErrorType Cmd)
 
 type TransmissionOrError = (Signature, SignedOrError)
 
-type RawTransmission = (ByteString, ByteString, ByteString)
+type RawTransmission = (ByteString, ByteString, ByteString, ByteString)
 
 data Command (a :: Party) where
   NEW :: RecipientKey -> Command Recipient
@@ -137,6 +137,8 @@ serializeCommand = \case
 
 type Encoded = ByteString
 
+type CorrelationId = ByteString
+
 type PublicKey = Encoded
 
 type Signature = Encoded

src/Simplex/Messaging/Transport.hs

@@ -80,20 +80,22 @@ getBytes :: MonadIO m => Handle -> Int -> m ByteString
 getBytes h = liftIO . B.hGet h
 
 tPutRaw :: MonadIO m => Handle -> RawTransmission -> m ()
-tPutRaw h (signature, queueId, command) = do
-  putLn h (encode signature)
-  putLn h (encode queueId)
+tPutRaw h (signature, corrId, queueId, command) = do
+  putLn h signature
+  putLn h corrId
+  putLn h queueId
   putLn h command
 
-tGetRaw :: MonadIO m => Handle -> m (Either String RawTransmission)
+tGetRaw :: MonadIO m => Handle -> m RawTransmission
 tGetRaw h = do
-  signature <- decode <$> getLn h
-  queueId <- decode <$> getLn h
+  signature <- getLn h
+  corrId <- getLn h
+  queueId <- getLn h
   command <- getLn h
-  return $ liftM2 (,,command) signature queueId
+  return (signature, corrId, queueId, command)
 
 tPut :: MonadIO m => Handle -> Transmission -> m ()
-tPut h (signature, (queueId, command)) = tPutRaw h (signature, queueId, serializeCommand command)
+tPut h (signature, (corrId, queueId, command)) = tPutRaw h (encode signature, corrId, encode queueId, serializeCommand command)
 
 fromClient :: Cmd -> Either ErrorType Cmd
 fromClient = \case
@@ -108,19 +110,22 @@ fromServer = \case
 -- | get client and server transmissions
 -- `fromParty` is used to limit allowed senders - `fromClient` or `fromServer` should be used
 tGet :: forall m. MonadIO m => (Cmd -> Either ErrorType Cmd) -> Handle -> m TransmissionOrError
-tGet fromParty h = tGetRaw h >>= either (const tError) tParseLoadBody
+tGet fromParty h = do
+  (signature, corrId, queueId, command) <- tGetRaw h
+  let decodedTransmission = liftM2 (,corrId,,command) (decode signature) (decode queueId)
+  either (const $ tError corrId) tParseLoadBody decodedTransmission
   where
-    tError :: m TransmissionOrError
-    tError = return (B.empty, (B.empty, Left $ SYNTAX errBadTransmission))
+    tError :: ByteString -> m TransmissionOrError
+    tError corrId = return (B.empty, (corrId, B.empty, Left $ SYNTAX errBadTransmission))
 
     tParseLoadBody :: RawTransmission -> m TransmissionOrError
-    tParseLoadBody t@(signature, queueId, command) = do
+    tParseLoadBody t@(signature, corrId, queueId, command) = do
       let cmd = parseCommand command >>= fromParty >>= tCredentials t
       fullCmd <- either (return . Left) cmdWithMsgBody cmd
-      return (signature, (queueId, fullCmd))
+      return (signature, (corrId, queueId, fullCmd))
 
     tCredentials :: RawTransmission -> Cmd -> Either ErrorType Cmd
-    tCredentials (signature, queueId, _) cmd = case cmd of
+    tCredentials (signature, _, queueId, _) cmd = case cmd of
       -- IDS response should not have queue ID
       Cmd SBroker (IDS _ _) -> Right cmd
       -- ERROR response does not always have queue ID