diff --git a/src/ConnStore.hs b/src/ConnStore.hs index e6825de54..711eb8b5e 100644 --- a/src/ConnStore.hs +++ b/src/ConnStore.hs @@ -1,27 +1,13 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} --- {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} module ConnStore where -import Control.Concurrent.STM -import Data.Map (Map) -import qualified Data.Map as M -import Polysemy -import Polysemy.Input +import Data.Singletons import Transmission -type SMPResult a = Either SMPError a - -data SMPError = CmdError | SyntaxError | AuthError | InternalError +type SMPResult a = Either ErrorType a data Connection = Connection { recipientId :: ConnId, @@ -31,23 +17,13 @@ data Connection = Connection active :: Bool } -data ConnStore m a where - CreateConn :: RecipientKey -> ConnStore m (SMPResult Connection) - GetConn :: Party -> ConnId -> ConnStore m (SMPResult Connection) +class MonadConnStore s m where + createConn :: s -> RecipientKey -> m (SMPResult Connection) + getConn :: s -> Sing (a :: Party) -> ConnId -> m (SMPResult Connection) --- SecureConn :: RecipientId -> SenderKey -> ConnStore m (SMPResult ()) --- SuspendConn :: RecipientId -> ConnStore m (SMPResult ()) --- DeleteConn :: RecipientId -> ConnStore m (SMPResult ()) - -makeSem ''ConnStore - -data ConnStoreData = ConnStoreData - { connections :: Map RecipientId Connection, - senders :: Map SenderId RecipientId - } - -newConnStore :: STM (TVar ConnStoreData) -newConnStore = newTVar ConnStoreData {connections = M.empty, senders = M.empty} +-- secureConn :: RecipientId -> SenderKey -> m (SMPResult ()) +-- suspendConn :: RecipientId -> m (SMPResult ()) +-- deleteConn :: RecipientId -> m (SMPResult ()) newConnection :: RecipientKey -> Connection newConnection rKey = @@ -58,28 +34,3 @@ newConnection rKey = senderKey = Nothing, active = True } - -runConnStoreSTM :: Member (Embed STM) r => Sem (ConnStore ': r) a -> Sem (Input (TVar ConnStoreData) ': r) a -runConnStoreSTM = reinterpret $ \case - CreateConn rKey -> do - store <- input - db <- embed $ readTVar store - let conn@Connection {senderId, recipientId} = newConnection rKey - db' = - ConnStoreData - { connections = M.insert recipientId conn (connections db), - senders = M.insert senderId recipientId (senders db) - } - embed $ writeTVar store db' - return $ Right conn - GetConn Recipient rId -> do - db <- input >>= embed . readTVar - return $ getRcpConn db rId - GetConn Sender sId -> do - db <- input >>= embed . readTVar - return $ maybeError (getRcpConn db) $ M.lookup sId $ senders db - GetConn Broker _ -> do - return $ Left InternalError - where - maybeError = maybe (Left AuthError) - getRcpConn db rId = maybeError Right $ M.lookup rId $ connections db diff --git a/src/ConnStore/STM.hs b/src/ConnStore/STM.hs new file mode 100644 index 000000000..11941ad7b --- /dev/null +++ b/src/ConnStore/STM.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} + +module ConnStore.STM where + +import ConnStore +import Control.Monad.IO.Unlift +import Data.Map (Map) +import qualified Data.Map as M +import Transmission +import UnliftIO.STM + +data ConnStoreData = ConnStoreData + { connections :: Map RecipientId Connection, + senders :: Map SenderId RecipientId + } + +type STMConnStore = TVar ConnStoreData + +newConnStore :: STM STMConnStore +newConnStore = newTVar ConnStoreData {connections = M.empty, senders = M.empty} + +instance MonadUnliftIO m => MonadConnStore STMConnStore m where + createConn store rKey = atomically $ do + db <- readTVar store + let conn@Connection {senderId, recipientId} = newConnection rKey + db' = + ConnStoreData + { connections = M.insert recipientId conn (connections db), + senders = M.insert senderId recipientId (senders db) + } + writeTVar store db' + return $ Right conn + getConn store SRecipient rId = atomically $ do + db <- readTVar store + return $ getRcpConn db rId + getConn store SSender sId = atomically $ do + db <- readTVar store + return $ maybeAuth (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) + +getRcpConn :: ConnStoreData -> RecipientId -> Either ErrorType Connection +getRcpConn db rId = maybeAuth Right $ M.lookup rId $ connections db diff --git a/src/EnvStm.hs b/src/Env/STM.hs similarity index 95% rename from src/EnvStm.hs rename to src/Env/STM.hs index 76bc8fe06..e8ce02ba1 100644 --- a/src/EnvStm.hs +++ b/src/Env/STM.hs @@ -1,9 +1,9 @@ {-# LANGUAGE DuplicateRecordFields #-} {-# LANGUAGE NamedFieldPuns #-} -module EnvSTM where +module Env.STM where -import ConnStore +import ConnStore.STM import Control.Concurrent.STM import qualified Data.Map as M import qualified Data.Set as S diff --git a/src/Main.hs b/src/Main.hs index 5f4b76aa2..10c0a1c21 100644 --- a/src/Main.hs +++ b/src/Main.hs @@ -3,6 +3,7 @@ {-# LANGUAGE DuplicateRecordFields #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -14,7 +15,7 @@ import Control.Monad import Control.Monad.IO.Unlift import Control.Monad.Reader import qualified Data.ByteString.Char8 as B -import EnvSTM +import Env.STM import Network.Socket import Text.Read import Transmission @@ -41,12 +42,12 @@ runTCPServer server = putLn h "Welcome" forkFinally (server h) (const $ hClose h) -runClient :: MonadUnliftIO m => Handle -> m () +runClient :: (MonadUnliftIO m, MonadReader Env m) => Handle -> m () runClient h = do c <- atomically $ newClient h void $ race (client c) (receive c) -receive :: MonadUnliftIO m => Client -> m () +receive :: (MonadUnliftIO m, MonadReader Env m) => Client -> m () receive Client {handle, channel} = forever $ do signature <- getLn handle connId <- getLn handle @@ -54,43 +55,68 @@ receive Client {handle, channel} = forever $ do cmdOrError <- parseReadVerifyTransmission handle signature connId command atomically $ writeTChan channel cmdOrError -parseReadVerifyTransmission :: MonadUnliftIO m => Handle -> String -> String -> String -> m SomeSigned +parseReadVerifyTransmission :: forall m. (MonadUnliftIO m, MonadReader Env m) => Handle -> String -> String -> String -> m SomeSigned parseReadVerifyTransmission h signature connId command = do let cmd = parseCommand command cmd' <- case cmd of Cmd SBroker _ -> return cmd - Cmd _ (CREATE _) -> + Cmd _ (CREATE _) -> signed False cmd errHasCredentials + Cmd _ (SEND msgBody) -> getSendMsgBody msgBody + Cmd _ _ -> verifyConnSignature cmd -- signed True cmd errNoCredentials + return (Just connId, cmd') + where + signed :: Bool -> Cmd -> Int -> m Cmd + signed isSigned cmd errCode = return - if signature == "" && connId == "" + if isSigned == (signature /= "") && isSigned == (connId /= "") then cmd - else smpError SYNTAX - Cmd _ (SEND msgBody) -> + else syntaxError errCode + getSendMsgBody :: MsgBody -> m Cmd + getSendMsgBody msgBody = if connId == "" - then return $ smpError SYNTAX + then return $ syntaxError errNoConnectionId else case B.unpack msgBody of ':' : body -> return . smpSend $ B.pack body sizeStr -> case readMaybe sizeStr :: Maybe Int of Just size -> do body <- getBytes h size s <- getLn h - return if s == "" then smpSend body else smpError SYNTAX - Nothing -> return $ smpError SYNTAX - Cmd _ _ -> - return - if signature == "" || connId == "" - then smpError SYNTAX - else cmd - return (Just connId, cmd') + return if s == "" then smpSend body else syntaxError errMessageBodySize + Nothing -> return $ syntaxError errMessageBody + verifyConnSignature :: Cmd -> m Cmd + verifyConnSignature cmd@(Cmd party _) = + if null signature || null connId + then return $ syntaxError errNoCredentials + else do + store <- asks connStore + getConn store party connId >>= \case + Right Connection {recipientKey, senderKey} -> do + res <- case party of + SRecipient -> verifySignature recipientKey + SSender -> case senderKey of + Just key -> verifySignature key + Nothing -> return False + SBroker -> return False + if res then return cmd else return $ smpError AUTH + Left err -> return $ smpError err + verifySignature :: Encoded -> m Bool + verifySignature key = return $ signature == key -client :: MonadIO m => Client -> m () +client :: (MonadUnliftIO m, MonadReader Env m) => Client -> m () client Client {handle, channel} = loop where loop = forever $ do (_, cmdOrErr) <- atomically $ readTChan channel - let response = case cmdOrErr of - Cmd SRecipient _ -> "OK" - Cmd SSender _ -> "OK" - Cmd SBroker (ERROR t) -> "ERROR " ++ show t - _ -> "ERROR INTERNAL" + response <- case cmdOrErr of + Cmd SRecipient (CREATE recipientKey) -> do + store <- asks connStore + conn <- createConn store recipientKey + case conn of + Right Connection {recipientId, senderId} -> return $ "CONN " ++ recipientId ++ " " ++ senderId + Left e -> return $ "ERROR " ++ show e + Cmd SRecipient _ -> return "OK" + Cmd SSender _ -> return "OK" + Cmd SBroker (ERROR e) -> return $ "ERROR " ++ show e + _ -> return "ERROR INTERNAL" putLn handle response liftIO $ print cmdOrErr diff --git a/src/Transmission.hs b/src/Transmission.hs index 318e38d7f..1e9fa925c 100644 --- a/src/Transmission.hs +++ b/src/Transmission.hs @@ -56,19 +56,23 @@ parseCommand command = case words command of ["SUSPEND"] -> rCmd SUSPEND ["DELETE"] -> rCmd DELETE ["SEND", msgBody] -> smpSend $ B.pack msgBody - "CREATE" : _ -> smpError SYNTAX - "SUB" : _ -> smpError SYNTAX - "SECURE" : _ -> smpError SYNTAX - "DELMSG" : _ -> smpError SYNTAX - "SUSPEND" : _ -> smpError SYNTAX - "DELETE" : _ -> smpError SYNTAX - "SEND" : _ -> smpError SYNTAX - _ -> smpError CMD + "CREATE" : _ -> err + "SUB" : _ -> err + "SECURE" : _ -> err + "DELMSG" : _ -> err + "SUSPEND" : _ -> err + "DELETE" : _ -> err + "SEND" : _ -> err + _ -> syntaxError errUnknownCommand where + err = syntaxError errBadParameters rCmd = Cmd SRecipient +syntaxError :: Int -> Cmd +syntaxError err = smpError $ SYNTAX err + smpError :: ErrorType -> Cmd -smpError = Cmd SBroker . ERROR +smpError errType = Cmd SBroker $ ERROR errType smpSend :: MsgBody -> Cmd smpSend = Cmd SSender . SEND @@ -95,4 +99,25 @@ type Timestamp = Encoded type MsgBody = B.ByteString -data ErrorType = CMD | SYNTAX | AUTH | INTERNAL deriving (Show) +data ErrorType = SYNTAX Int | AUTH | INTERNAL deriving (Show) + +errUnknownCommand :: Int +errUnknownCommand = 1 + +errBadParameters :: Int +errBadParameters = 2 + +errNoCredentials :: Int +errNoCredentials = 3 + +errHasCredentials :: Int +errHasCredentials = 4 + +errNoConnectionId :: Int +errNoConnectionId = 5 + +errMessageBody :: Int +errMessageBody = 6 + +errMessageBodySize :: Int +errMessageBodySize = 7 diff --git a/src/Transport.hs b/src/Transport.hs index 244b6f8d2..c9230cb25 100644 --- a/src/Transport.hs +++ b/src/Transport.hs @@ -1,12 +1,12 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE FlexibleContexts #-} module Transport where import Control.Monad.IO.Class -import Control.Monad.IO.Unlift import Control.Monad.Reader import qualified Data.ByteString.Char8 as B -import EnvSTM +import Env.STM import Network.Socket import System.IO @@ -39,5 +39,5 @@ putLn h = liftIO . hPutStrLn h getLn :: MonadIO m => Handle -> m String getLn = liftIO . hGetLine -getBytes :: MonadUnliftIO m => Handle -> Int -> m B.ByteString +getBytes :: MonadIO m => Handle -> Int -> m B.ByteString getBytes h = liftIO . B.hGet h