SECURE command, tests

src/ConnStore.hs

@@ -7,8 +7,6 @@ module ConnStore where
 import Data.Singletons
 import Transmission
 
-type SMPResult a = Either ErrorType a
-
 data Connection = Connection
   { recipientId :: ConnId,
     recipientKey :: PublicKey,
@@ -18,12 +16,12 @@ data Connection = Connection
   }
 
 class MonadConnStore s m where
-  createConn :: s -> RecipientKey -> m (SMPResult Connection)
-  getConn :: s -> Sing (a :: Party) -> ConnId -> m (SMPResult Connection)
+  createConn :: s -> RecipientKey -> m (Either ErrorType Connection)
+  getConn :: s -> Sing (a :: Party) -> ConnId -> m (Either ErrorType Connection)
+  secureConn :: s -> RecipientId -> SenderKey -> m (Either ErrorType ())
 
--- secureConn :: RecipientId -> SenderKey -> m (SMPResult ())
--- suspendConn :: RecipientId -> m (SMPResult ())
--- deleteConn :: RecipientId -> m (SMPResult ())
+-- suspendConn :: RecipientId -> m (Either ErrorType ())
+-- deleteConn :: RecipientId -> m (Either ErrorType ())
 
 newConnection :: RecipientKey -> Connection
 newConnection rKey =

src/ConnStore/STM.hs

@@ -1,7 +1,7 @@
+{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE UndecidableInstances #-}
@@ -26,27 +26,38 @@ newConnStore :: STM STMConnStore
 newConnStore = newTVar ConnStoreData {connections = M.empty, senders = M.empty}
 
 instance MonadUnliftIO m => MonadConnStore STMConnStore m where
-  createConn store rKey = atomically $ do
+  createConn store rKey = atomically do
     db <- readTVar store
-    let conn@Connection {senderId, recipientId} = newConnection rKey
+    let c@Connection {recipientId = rId, senderId = sId} = newConnection rKey
         db' =
           ConnStoreData
-            { connections = M.insert recipientId conn (connections db),
-              senders = M.insert senderId recipientId (senders db)
+            { connections = M.insert rId c (connections db),
+              senders = M.insert sId rId (senders db)
             }
     writeTVar store db'
-    return $ Right conn
-  getConn store SRecipient rId = atomically $ do
+    return $ Right c
+
+  -- TODO do not return suspended connections
+  getConn store SRecipient rId = atomically do
     db <- readTVar store
     return $ getRcpConn db rId
-  getConn store SSender sId = atomically $ do
+  getConn store SSender sId = atomically do
     db <- readTVar store
-    return $ maybeAuth (getRcpConn db) $ M.lookup sId $ senders db
-  getConn _ SBroker _ = atomically $ do
+    return $ maybe (Left AUTH) (getRcpConn db) $ M.lookup sId $ senders db
+  getConn _ SBroker _ = atomically do
     return $ Left INTERNAL
 
-maybeAuth :: (a -> Either ErrorType b) -> Maybe a -> Either ErrorType b
-maybeAuth = maybe (Left AUTH)
+  secureConn store rId sKey = atomically do
+    db <- readTVar store
+    let conn = getRcpConn db rId
+    either (return . Left) (updateConn db) conn
+    where
+      updateConn db c = case senderKey c of
+        Just _ -> return $ Left AUTH
+        Nothing -> do
+          let db' = db {connections = M.insert rId c {senderKey = Just sKey} (connections db)}
+          writeTVar store db'
+          return $ Right ()
 
 getRcpConn :: ConnStoreData -> RecipientId -> Either ErrorType Connection
-getRcpConn db rId = maybeAuth Right $ M.lookup rId $ connections db
+getRcpConn db rId = maybe (Left AUTH) Right . M.lookup rId $ connections db

src/Server.hs

@@ -26,12 +26,12 @@ import qualified UnliftIO.Exception as E
 import UnliftIO.IO
 import UnliftIO.STM
 
-runSMPServer :: ServiceName -> Natural -> IO ()
+runSMPServer :: MonadUnliftIO m => ServiceName -> Natural -> m ()
 runSMPServer port queueSize = do
   env <- atomically $ newEnv port queueSize
   runReaderT (runTCPServer runClient) env
 
-runTCPServer :: (MonadReader Env m, MonadUnliftIO m) => (Handle -> m ()) -> m ()
+runTCPServer :: (MonadUnliftIO m, MonadReader Env m) => (Handle -> m ()) -> m ()
 runTCPServer server =
   E.bracket startTCPServer (liftIO . close) $ \sock -> forever $ do
     h <- acceptTCPConn sock
@@ -58,7 +58,7 @@ receive h Client {queue} = forever $ do
 verifyTransmission :: forall m. (MonadUnliftIO m, MonadReader Env m) => Signature -> ConnId -> Cmd -> m Signed
 verifyTransmission signature connId cmd = do
   (connId,) <$> case cmd of
-    Cmd SBroker _ -> return $ smpErr INTERNAL
+    Cmd SBroker _ -> return $ smpErr INTERNAL -- it can only be client command, because `fromClient` was used
     Cmd SRecipient (CREATE _) -> return cmd
     Cmd SRecipient _ -> withConnection SRecipient $ verifySignature . recipientKey
     Cmd SSender (SEND _) -> withConnection SSender $ verifySend . senderKey
@@ -80,18 +80,35 @@ verifyTransmission signature connId cmd = do
     smpErr e = Cmd SBroker $ ERROR e
     authErr = smpErr AUTH
 
-client :: (MonadUnliftIO m, MonadReader Env m) => Handle -> Client -> m ()
+client :: forall m. (MonadUnliftIO m, MonadReader Env m) => Handle -> Client -> m ()
 client h Client {queue} = loop
   where
     loop = forever $ do
       (connId, cmd) <- atomically $ readTBQueue queue
-      response <- case cmd of
-        Cmd SRecipient (CREATE recipientKey) -> do
-          store <- asks connStore
-          conn <- createConn store recipientKey
-          return . Cmd SBroker $ case conn of
-            Right Connection {recipientId, senderId} -> CONN recipientId senderId
-            Left e -> ERROR e
-        Cmd SBroker _ -> return cmd
-        Cmd _ _ -> return $ Cmd SBroker OK
-      tPut h ("", (connId, response)) -- empty signature
+      signed <- processCommand connId cmd
+      tPut h ("", signed)
+
+    processCommand :: ConnId -> Cmd -> m Signed
+    processCommand connId cmd = do
+      st <- asks connStore
+      case cmd of
+        Cmd SRecipient (CREATE recipientKey) ->
+          either (mkSigned "" . ERROR) connResponce
+            <$> createConn st recipientKey
+        Cmd SRecipient SUB -> do
+          -- TODO message subscription
+          return ok
+        Cmd SRecipient (SECURE senderKey) -> do
+          mkSigned connId . either ERROR (const OK)
+            <$> secureConn st connId senderKey
+        Cmd SBroker _ -> return (connId, cmd)
+        Cmd _ _ -> return ok
+      where
+        ok :: Signed
+        ok = (connId, Cmd SBroker OK)
+
+        mkSigned :: ConnId -> Command 'Broker -> Signed
+        mkSigned cId command = (cId, Cmd SBroker command)
+
+        connResponce :: Connection -> Signed
+        connResponce Connection {recipientId = rId, senderId = sId} = mkSigned rId $ CONN rId sId