http://www.haskell.org/pipermail/beginners/2013-November/012912.html

iperf.hs

{-# LANGUAGE BangPatterns #-}
import Data.Time (DiffTime, utctDayTime, getCurrentTime)
import Foreign (ForeignPtr, Int64, Ptr, mallocForeignPtr, withForeignPtr)
import Foreign.Storable (sizeOf, peek, poke)
import GHC.IO.Handle ( Handle, BufferMode(NoBuffering)
                     , hClose, hGetBuf, hSetBuffering, hPutBuf)
import Network ( Socket, HostName, PortID(PortNumber), PortNumber, accept
               , connectTo, listenOn, withSocketsDo)
import System.Console.GetOpt ( ArgOrder(Permute), ArgDescr(..), OptDescr(..)
                             , getOpt, usageInfo)
import System.Environment (getArgs)

data Flag = Client String
          | Server
          deriving (Show)

-- Command line arguments for server and client mode.
options :: [OptDescr Flag]
options =
  [ Option ['c'] ["client"] (ReqArg Client "HOST") "Connect to host as client."
  , Option ['s'] ["server"] (NoArg Server) "Run as server."
  ]

failWithMessage :: String -> IO a
failWithMessage = ioError . userError . (++ usageInfo "" options)

printWords :: [String] -> IO ()
printWords = putStrLn . unwords

perfOpts :: [String] -> IO [Flag]
perfOpts argv =
  case getOpt Permute options argv of
    ([], [], []) -> failWithMessage "At least one Option is needed."
    (o, [], [])  -> return o
    (_, _, errs) -> failWithMessage (concat errs)

-- Some constants

defaultPort :: PortID
defaultPort = PortNumber 8456

type Trans = Int64

numBytes :: Int
numBytes = sizeOf (0 :: Trans)

initial, intermediate, final :: Trans
initial = -128
intermediate = 0
final = 127

makePtr :: IO (ForeignPtr Trans)
makePtr = mallocForeignPtr

-- The client connects to the server, initializes the connection
-- by sending initial, sends then for the amount of at least
-- 10 secs as much intermediate (0's) as possible and finally
-- sends final.

runClient :: String -> IO ()
runClient host = do
  printWords ["Running as client connecting to", host]
  hdl <- initClient host defaultPort
  ptr <- makePtr
  fillInitial ptr
  withForeignPtr ptr (sendBuf hdl)
  fillIntermediate ptr
  curTime <- fmap utctDayTime getCurrentTime
  (sentBytes, time) <- sendIntermediateAndFinal ptr hdl curTime 10 0
  printWords ["Sent", show sentBytes, "Bytes in", show time, "seconds."]
  hClose hdl

initClient :: HostName -> PortID -> IO Handle
initClient host port = withSocketsDo $ do
  hdl <- connectTo host port
  hSetBuffering hdl NoBuffering
  return hdl

sendIntermediateAndFinal :: ForeignPtr Trans -> Handle -> DiffTime
                         -> DiffTime -> Int -> IO (Int, DiffTime)
sendIntermediateAndFinal ptr hdl start duration !sent = do
  curTime <- fmap utctDayTime getCurrentTime
  if (curTime - start) > duration
    then do fillFinal ptr
            withForeignPtr ptr (sendBuf hdl)
            finishedTime <- fmap utctDayTime getCurrentTime
            return (sent + numBytes, finishedTime - start)
    else do withForeignPtr ptr (sendBuf hdl)
            sendIntermediateAndFinal ptr hdl start duration (sent + numBytes)

sendBuf :: Handle -> Ptr Trans -> IO ()
sendBuf hdl buf = hPutBuf hdl buf numBytes

fillPtr :: Trans -> ForeignPtr Trans -> IO ()
fillPtr num ptr = withForeignPtr ptr (\p -> poke p num)

fillInitial, fillIntermediate, fillFinal :: ForeignPtr Trans -> IO ()
fillInitial = fillPtr initial
fillIntermediate = fillPtr intermediate
fillFinal = fillPtr final

-- The server simply accepts connections, receives what it can get
-- and adds up the received bytes and transmission time.

runServer :: IO ()
runServer = do
  putStrLn "Running as server."
  initServer defaultPort >>= handleConnection

initServer :: PortID -> IO Socket
initServer port = withSocketsDo $ listenOn port

handleConnection :: Socket -> IO ()
handleConnection socket = do
  (hdl, host, _port) <- acceptConnection socket
  printWords ["Connection from", host]
  ptr <- makePtr
  (num, _bytes) <- receive ptr hdl
  curTime <- fmap utctDayTime getCurrentTime
  if num == initial
    then do (received, time) <- receiveData ptr hdl curTime 0
            printWords ["Received", show received, "in", show time, "seconds."]
            hClose hdl
            handleConnection socket
    else do printWords ["ERR: Expected", show initial, "got", show num]
            hClose hdl
            handleConnection socket

acceptConnection :: Socket -> IO (Handle, HostName, PortNumber)
acceptConnection socket = do
  (hdl, host, port) <- accept socket
  hSetBuffering hdl NoBuffering
  return (hdl, host, port)

receiveData :: ForeignPtr Trans -> Handle -> DiffTime -> Int
            -> IO (Int, DiffTime)
receiveData ptr hdl started !received = do
  (num, bytes) <- receive ptr hdl
  if num /= final
    then receiveData ptr hdl started (received + bytes)
    else do curTime <- fmap utctDayTime getCurrentTime
            return (received + bytes, curTime - started)

receive :: ForeignPtr Trans -> Handle -> IO (Trans, Int)
receive ptr hdl = do
  withForeignPtr ptr (\p -> hGetBuf hdl p numBytes)
  num <- withForeignPtr ptr peek
  return (num, numBytes)

-- The main routine decides whether to run as server or client.

main :: IO ()
main = do
  opts <- getArgs >>= perfOpts
  case opts of
    [Server]      -> runServer
    [Client host] -> runClient host
    _             -> failWithMessage "Wrong number of arguments!"