Simple demo UDP server and client

Test.hs

module Main where

import Control.Applicative
import Control.Concurrent
import Control.Error
import Control.Monad
import Control.Monad.Trans (lift)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Binary.Get as B
import qualified Data.Binary.Put as B
import qualified Data.Binary.IEEE754 as F
import Data.Time
import Network.Socket
import qualified Network.Socket.ByteString as NBS
import System.Environment

    
------------------------------------------------------------------------------
-- Data generator (client)
runClient :: Socket -> SockAddr -> IO ()
runClient sock sockAddr = forever $ do
  putStrLn "Type something, we'll UDP it out"
  l <- getLine
  nBytes <- sendTo sock l sockAddr
  putStrLn $ "Sent " ++ show nBytes ++ " bytes of " ++ l

runJonClient :: Socket -> SockAddr -> IO ()
runJonClient sock sockAddr = do
  putStrLn "Sending floats"
  let bs =  B.runPut $ forM_ [1,2,3,4,5 :: Float] F.putFloat32le
  nBytes <- NBS.sendTo sock (BSL.toStrict bs) sockAddr
  putStrLn $ "Sent " ++ show nBytes ++ " bytes"

------------------------------------------------------------------------------
-- Data listener (server)
runServer :: Socket -> SockAddr -> IO ()
runServer sock sockAddr = do
  bind sock sockAddr
  forever $ do
    putStrLn "Awaiting message"
    (msg, _, _) <- recvFrom sock 4492
    putStrLn $ "Received " ++ msg

------------------------------------------------------------------------------
-- Float Listener (server)
runJonServer :: Socket -> SockAddr -> IO ()
runJonServer sock sockAddr = do
  bind sock sockAddr
  putStrLn "Awaiting messages."
  forever $ do
    (msg, _) <- NBS.recvFrom sock 10000
    let nFloats = BS.length msg `div` 4 :: Int
        xs      = B.runGet ((replicateM nFloats F.getFloat32le)) (BSL.fromStrict msg :: BSL.ByteString)
    putStrLn $ "Received: " ++ show (BS.length msg) ++ " bytes. Decoded to: " ++ show xs


runTimingServer :: Socket -> SockAddr -> IO ()
runTimingServer sock sockAddr =
  bind sock sockAddr >> go 0 0 0 0
  where
    go :: Double -> Double -> Integer -> Integer -> IO ()
    go sX sXX nGood nBad = do
      (msg, _, _) <- recvFrom sock 4082
      t <- getCurrentTime
      let (sX', sXX', nGood', nBad') = case readMay msg of
            Just msgTime ->
              let dt = realToFrac (diffUTCTime t msgTime)
              in  (sX + dt, sXX + dt^(2::Int), nGood+1, nBad)
            Nothing ->
              (sX, sXX, nGood, nBad + 1)
          meanX = sX / (fromIntegral nGood)
          sdX   = sqrt ( (sXX/ fromIntegral nGood) - (meanX^(2::Int)) )
      when ((nGood' + nBad') `mod` 100 == 0) $
        putStrLn $ unwords ["Mean:",  show meanX
                           ,"StDev:", show sdX
                           ,"NGood:", show nGood
                           ,"NBad:",  show nBad
                           ]
      go sX' sXX' nGood' nBad'


runTimingClient :: Socket -> SockAddr -> IO ()
runTimingClient sock sockAddr = forever $ do
  t <- getCurrentTime
  sendTo sock (show t) sockAddr
  threadDelay 500
  
------------------------------------------------------------------------------
main :: IO ()
main = do
  args  <- getArgs
  case args of
    [mode,host,port] -> do
      sock     <- socket AF_INET Datagram defaultProtocol
      sockAddr <- runEitherT $ getSockAddr host port
      case (sockAddr,mode) of
        (Left  e,  _         )       -> error $
                                        "Socket address name error: " ++ e
        (Right sa, "listener")      -> runJonServer sock sa
        (Right sa, "timedListener") -> runTimingServer sock sa
        (Right sa, "talker"  )      -> runJonClient sock sa
        (Right sa, "timedTalker")   -> runTimingClient sock sa
        _          -> usage
    _ -> usage
 where usage = error "Usage: test-udp [talker|listener] hostIP port"


------------------------------------------------------------------------------
getSockAddr :: String -> String -> EitherT String IO SockAddr
getSockAddr host port = SockAddrInet
       <$> (toEnum <$> readMay port ?? "Couldn't read port number")
       <*> lift (inet_addr host)