move SMP server from agent commands NEW/JOIN to agent config (#129)

* move SMP server from agent commands NEW/JOIN to agent config

* fix SMPServer parser

* update agent protocol - server management

* enable agent test

* agent test with 2 servers

* create reply queue on the configured server

* choose random server

* swap bind

src/Simplex/Messaging/Agent.hs

@@ -16,6 +16,7 @@ module Simplex.Messaging.Agent
   )
 where
 
+import Control.Concurrent.STM (stateTVar)
 import Control.Logger.Simple (logInfo, showText)
 import Control.Monad.Except
 import Control.Monad.IO.Unlift (MonadUnliftIO)
@@ -23,6 +24,8 @@ import Control.Monad.Reader
 import Crypto.Random (MonadRandom)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import Data.List.NonEmpty (NonEmpty (..))
+import qualified Data.List.NonEmpty as L
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeUtf8)
 import Data.Time.Clock
@@ -39,6 +42,7 @@ import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Transport (putLn, runTCPServer)
 import Simplex.Messaging.Util (bshow)
 import System.IO (Handle)
+import System.Random (randomR)
 import UnliftIO.Async (race_)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
@@ -51,7 +55,7 @@ runSMPAgentBlocking started cfg@AgentConfig {tcpPort} = runReaderT smpAgent =<<
   where
     smpAgent :: (MonadUnliftIO m', MonadReader Env m') => m' ()
     smpAgent = runTCPServer started tcpPort $ \h -> do
-      liftIO $ putLn h "Welcome to SMP v0.2.0 agent"
+      liftIO $ putLn h "Welcome to SMP v0.3.0 agent"
       c <- getSMPAgentClient
       logConnection c True
       race_ (connectClient h c) (runSMPAgentClient c)
@@ -127,7 +131,7 @@ withStore action = do
 processCommand :: forall m. AgentMonad m => AgentClient -> SQLiteStore -> ATransmission 'Client -> m ()
 processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
   case cmd of
-    NEW smpServer -> createNewConnection smpServer
+    NEW -> createNewConnection
     JOIN smpQueueInfo replyMode -> joinConnection smpQueueInfo replyMode
     SUB -> subscribeConnection connAlias
     SUBALL -> subscribeAll
@@ -135,25 +139,32 @@ processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
     OFF -> suspendConnection
     DEL -> deleteConnection
   where
-    createNewConnection :: SMPServer -> m ()
-    createNewConnection server = do
+    createNewConnection :: m ()
+    createNewConnection = do
       -- TODO create connection alias if not passed
       -- make connAlias Maybe?
-      (rq, qInfo) <- newReceiveQueue c server connAlias
+      srv <- getSMPServer
+      (rq, qInfo) <- newReceiveQueue c srv connAlias
       withStore $ createRcvConn st rq
       respond $ INV qInfo
 
+    getSMPServer :: m SMPServer
+    getSMPServer =
+      asks (smpServers . config) >>= \case
+        srv :| [] -> pure srv
+        servers -> do
+          gen <- asks randomServer
+          i <- atomically . stateTVar gen $ randomR (0, L.length servers - 1)
+          pure $ servers L.!! i
+
     joinConnection :: SMPQueueInfo -> ReplyMode -> m ()
-    joinConnection qInfo@(SMPQueueInfo srv _ _) replyMode = do
+    joinConnection qInfo (ReplyMode replyMode) = do
       -- TODO create connection alias if not passed
       -- make connAlias Maybe?
       (sq, senderKey, verifyKey) <- newSendQueue qInfo connAlias
       withStore $ createSndConn st sq
       connectToSendQueue c st sq senderKey verifyKey
-      case replyMode of
-        ReplyOn -> sendReplyQInfo srv sq
-        ReplyVia srv' -> sendReplyQInfo srv' sq
-        ReplyOff -> return ()
+      when (replyMode == On) $ createReplyQueue sq
       respond CON
 
     subscribeConnection :: ConnAlias -> m ()
@@ -216,8 +227,9 @@ processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
           removeSubscription c connAlias
           delConn
 
-    sendReplyQInfo :: SMPServer -> SndQueue -> m ()
-    sendReplyQInfo srv sq = do
+    createReplyQueue :: SndQueue -> m ()
+    createReplyQueue sq = do
+      srv <- getSMPServer
       (rq, qInfo) <- newReceiveQueue c srv connAlias
       withStore $ upgradeSndConnToDuplex st connAlias rq
       senderTimestamp <- liftIO getCurrentTime

src/Simplex/Messaging/Agent/Env/SQLite.hs

@@ -7,14 +7,18 @@ module Simplex.Messaging.Agent.Env.SQLite where
 
 import Control.Monad.IO.Unlift
 import Crypto.Random
+import Data.List.NonEmpty (NonEmpty)
 import Network.Socket
 import Numeric.Natural
 import Simplex.Messaging.Agent.Store.SQLite
+import Simplex.Messaging.Agent.Transmission (SMPServer)
 import Simplex.Messaging.Client
+import System.Random (StdGen, newStdGen)
 import UnliftIO.STM
 
 data AgentConfig = AgentConfig
   { tcpPort :: ServiceName,
+    smpServers :: NonEmpty SMPServer,
     rsaKeySize :: Int,
     connIdBytes :: Int,
     tbqSize :: Natural,
@@ -26,15 +30,17 @@ data Env = Env
   { config :: AgentConfig,
     idsDrg :: TVar ChaChaDRG,
     clientCounter :: TVar Int,
-    reservedMsgSize :: Int
+    reservedMsgSize :: Int,
+    randomServer :: TVar StdGen
   }
 
 newSMPAgentEnv :: (MonadUnliftIO m, MonadRandom m) => AgentConfig -> m Env
 newSMPAgentEnv config = do
-  idsDrg <- drgNew >>= newTVarIO
+  idsDrg <- newTVarIO =<< drgNew
   _ <- createSQLiteStore $ dbFile config
   clientCounter <- newTVarIO 0
-  return Env {config, idsDrg, clientCounter, reservedMsgSize}
+  randomServer <- newTVarIO =<< liftIO newStdGen
+  return Env {config, idsDrg, clientCounter, reservedMsgSize, randomServer}
   where
     -- 1st rsaKeySize is used by the RSA signature in each command,
     -- 2nd - by encrypted message body header

src/Simplex/Messaging/Agent/Transmission.hs

@@ -23,6 +23,7 @@ import qualified Data.ByteString.Char8 as B
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.Kind (Type)
+import Data.String (IsString (..))
 import Data.Time.Clock (UTCTime)
 import Data.Time.ISO8601
 import Data.Type.Equality
@@ -75,7 +76,7 @@ data ACmd = forall p. ACmd (SAParty p) (ACommand p)
 deriving instance Show ACmd
 
 data ACommand (p :: AParty) where
-  NEW :: SMPServer -> ACommand Client -- response INV
+  NEW :: ACommand Client -- response INV
   INV :: SMPQueueInfo -> ACommand Agent
   JOIN :: SMPQueueInfo -> ReplyMode -> ACommand Client -- response OK
   CON :: ACommand Agent -- notification that connection is established
@@ -170,7 +171,7 @@ agentMessageP =
     a_msg = do
       size :: Int <- A.decimal <* A.endOfLine
       A_MSG <$> A.take size <* A.endOfLine
-    ackMode = " NO_ACK" $> AckMode Off <|> pure (AckMode On)
+    ackMode = AckMode <$> (" NO_ACK" $> Off <|> pure On)
 
 smpQueueInfoP :: Parser SMPQueueInfo
 smpQueueInfoP =
@@ -179,7 +180,7 @@ smpQueueInfoP =
 smpServerP :: Parser SMPServer
 smpServerP = SMPServer <$> server <*> optional port <*> optional kHash
   where
-    server = B.unpack <$> A.takeTill (A.inClass ":# ")
+    server = B.unpack <$> A.takeWhile1 (A.notInClass ":# ")
     port = A.char ':' *> (B.unpack <$> A.takeWhile1 A.isDigit)
     kHash = C.KeyHash <$> (A.char '#' *> base64P)
 
@@ -207,18 +208,21 @@ data SMPServer = SMPServer
   }
   deriving (Eq, Ord, Show)
 
+instance IsString SMPServer where
+  fromString = parseString . parseAll $ smpServerP
+
 type ConnAlias = ByteString
 
 type OtherPartyId = Encoded
 
-data Mode = On | Off deriving (Eq, Show, Read)
+data OnOff = On | Off deriving (Eq, Show, Read)
 
-newtype AckMode = AckMode Mode deriving (Eq, Show)
+newtype AckMode = AckMode OnOff deriving (Eq, Show)
 
 data SMPQueueInfo = SMPQueueInfo SMPServer SMP.SenderId EncryptionKey
   deriving (Eq, Show)
 
-data ReplyMode = ReplyOff | ReplyOn | ReplyVia SMPServer deriving (Eq, Show)
+newtype ReplyMode = ReplyMode OnOff deriving (Eq, Show)
 
 type EncryptionKey = C.PublicKey
 
@@ -294,7 +298,7 @@ instance Arbitrary SMPAgentError where arbitrary = genericArbitraryU
 
 commandP :: Parser ACmd
 commandP =
-  "NEW " *> newCmd
+  "NEW" $> ACmd SClient NEW
     <|> "INV " *> invResp
     <|> "JOIN " *> joinCmd
     <|> "SUB" $> ACmd SClient SUB
@@ -309,7 +313,6 @@ commandP =
     <|> "CON" $> ACmd SAgent CON
     <|> "OK" $> ACmd SAgent OK
   where
-    newCmd = ACmd SClient . NEW <$> smpServerP
     invResp = ACmd SAgent . INV <$> smpQueueInfoP
     joinCmd = ACmd SClient <$> (JOIN <$> smpQueueInfoP <*> replyMode)
     sendCmd = ACmd SClient . SEND <$> A.takeByteString
@@ -321,10 +324,7 @@ commandP =
       senderMeta <- "S=" *> partyMeta A.decimal
       msgBody <- A.takeByteString
       return $ ACmd SAgent MSG {recipientMeta, brokerMeta, senderMeta, msgIntegrity, msgBody}
-    replyMode =
-      " NO_REPLY" $> ReplyOff
-        <|> A.space *> (ReplyVia <$> smpServerP)
-        <|> pure ReplyOn
+    replyMode = ReplyMode <$> (" NO_REPLY" $> Off <|> pure On)
     partyMeta idParser = (,) <$> idParser <* "," <*> tsISO8601P <* A.space
     agentError = ACmd SAgent . ERR <$> agentErrorTypeP
 
@@ -342,7 +342,7 @@ parseCommand = parse commandP $ CMD SYNTAX
 
 serializeCommand :: ACommand p -> ByteString
 serializeCommand = \case
-  NEW srv -> "NEW " <> serializeServer srv
+  NEW -> "NEW"
   INV qInfo -> "INV " <> serializeSmpQueueInfo qInfo
   JOIN qInfo rMode -> "JOIN " <> serializeSmpQueueInfo qInfo <> replyMode rMode
   SUB -> "SUB"
@@ -367,9 +367,8 @@ serializeCommand = \case
   where
     replyMode :: ReplyMode -> ByteString
     replyMode = \case
-      ReplyOff -> " NO_REPLY"
-      ReplyVia srv -> " " <> serializeServer srv
-      ReplyOn -> ""
+      ReplyMode Off -> " NO_REPLY"
+      ReplyMode On -> ""
     showTs :: UTCTime -> ByteString
     showTs = B.pack . formatISO8601Millis
 
@@ -433,7 +432,7 @@ tGet party h = liftIO (tGetRaw h) >>= tParseLoadBody
     tConnAlias :: ARawTransmission -> ACommand p -> Either AgentErrorType (ACommand p)
     tConnAlias (_, connAlias, _) cmd = case cmd of
       -- NEW and JOIN have optional connAlias
-      NEW _ -> Right cmd
+      NEW -> Right cmd
       JOIN _ _ -> Right cmd
       -- ERROR response does not always have connAlias
       ERR _ -> Right cmd

src/Simplex/Messaging/Client.hs

@@ -244,9 +244,8 @@ sendSMPCommand SMPClient {sndQ, sentCommands, clientCorrId, tcpTimeout} pKey qId
 
     getNextCorrId :: STM CorrId
     getNextCorrId = do
-      i <- (+ 1) <$> readTVar clientCorrId
-      writeTVar clientCorrId i
-      return . CorrId $ bshow i
+      i <- stateTVar clientCorrId $ \i -> (i, i + 1)
+      pure . CorrId $ bshow i
 
     signTransmission :: ByteString -> ExceptT SMPClientError IO SignedRawTransmission
     signTransmission t = case pKey of

src/Simplex/Messaging/Crypto.hs

@@ -80,7 +80,7 @@ import Data.X509
 import Database.SQLite.Simple.FromField (FromField (..))
 import Database.SQLite.Simple.ToField (ToField (..))
 import Network.Transport.Internal (decodeWord32, encodeWord32)
-import Simplex.Messaging.Parsers (base64P, blobFieldParser, parseAll)
+import Simplex.Messaging.Parsers (base64P, blobFieldParser, parseAll, parseString)
 import Simplex.Messaging.Util (liftEitherError, (<$?>))
 
 newtype PublicKey = PublicKey {rsaPublicKey :: R.PublicKey} deriving (Eq, Show)
@@ -111,9 +111,6 @@ instance IsString FullPrivateKey where
 instance IsString PublicKey where
   fromString = parseString (decode >=> decodePubKey)
 
-parseString :: (ByteString -> Either String a) -> (String -> a)
-parseString parse = either error id . parse . B.pack
-
 instance ToField SafePrivateKey where toField = toField . encodePrivKey
 
 instance ToField PublicKey where toField = toField . encodePubKey

src/Simplex/Messaging/Parsers.hs

@@ -50,6 +50,9 @@ parseRead2 = parseRead $ do
   w2 <- A.takeTill (== ' ')
   pure $ w1 <> " " <> w2
 
+parseString :: (ByteString -> Either String a) -> (String -> a)
+parseString p = either error id . p . B.pack
+
 blobFieldParser :: Typeable k => Parser k -> FieldParser k
 blobFieldParser p = \case
   f@(Field (SQLBlob b) _) ->