Example ZeroMQ Haskell

zeromq-secure-example.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns      #-}
import Control.Applicative
import Control.Monad
import Control.Concurrent
import Control.Exception as X
import Data.Traversable
import Data.Monoid
import Data.String
import Data.Foldable
import Data.Foldable
import System.IO
import System.Exit
import System.Environment
import System.ZMQ4.Monadic
import qualified Data.ByteString.Char8 as CS
import Data.Restricted

main :: IO ()
main = do
    args@[addr, fromString -> name, read -> nrClients] <- getArgs
    when (length args /= 3) $ do
        hPutStrLn stderr "usage: prompt <address> <username> <number of clients>"
        exitFailure
    (serverPub, serverPriv) <- curveKeyPair
    ps <- for [1..nrClients] $ \i ->
            do (p,q) <- curveKeyPair
               void $ forkIO $ runClient i addr serverPub p q
               pure p
    runServer serverPriv ps addr name

runClient :: Int -> String -> Restricted Div5 CS.ByteString -> Restricted Div5 CS.ByteString -> Restricted Div5 CS.ByteString -> IO ()
runClient i addr serverP p q =
    X.catch (runClient' i addr serverP p q) (\e -> print (e :: X.SomeException))

runClient' :: Int -> String -> Restricted Div5 CS.ByteString -> Restricted Div5 CS.ByteString -> Restricted Div5 CS.ByteString -> IO ()
runClient' i addr serverP p q =
    runZMQ $ do
        sub <- socket Pull
        setCurveServerKey TextFormat serverP sub
        setCurveSecretKey TextFormat q sub
        setCurvePublicKey TextFormat p sub
        -- subscribe sub ""
        connect sub addr
        forever $ do
            receive sub >>= liftIO . putStrLn . (((show i ++ ": ") ++) . CS.unpack)
            liftIO $ hFlush stdout

runServer :: Restricted Div5 CS.ByteString -> [Restricted Div5 CS.ByteString] -> String -> CS.ByteString -> IO ()
runServer q ps addr name = runZMQ $ do
    pub <- socket Push
    setCurveServer True pub
    setCurveSecretKey TextFormat q pub
    traverse_ (\p -> setCurvePublicKey TextFormat p pub) ps
    bind pub "tcp://*:2000" -- addr
    forever $ do
       line <- liftIO $ fromString <$> getLine
       replicateM_ 10000 $ send pub [] ("from " <> name <> ": " <> line)