Haskell TCP server/client

Client.hs

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Concurrent
import Control.Exception
import Control.Monad
import GHC.IO.Exception
import Network.Socket
import System.IO
import System.IO.Error
import Util

main :: IO ()
main = runTCPClient "localhost" "5000" ping

ping :: Handle -> IO ()
ping h = do
  putStrLn "connected to server"
  forever $ do
    putLn h "ping"
    threadDelay 1000000
    line <- getLn h
    print line

runTCPClient :: HostName -> ServiceName -> (Handle -> IO a) -> IO a
runTCPClient host port client = do
  h <- startTCPClient host port
  client h `finally` hClose h

startTCPClient :: HostName -> ServiceName -> IO Handle
startTCPClient host port =
  withSocketsDo $
    resolve >>= foldM tryOpen (Left err) >>= either throwIO return
  where
    err :: IOException
    err = mkIOError NoSuchThing "no address" Nothing Nothing

    resolve :: IO [AddrInfo]
    resolve = do
      let hints = defaultHints {addrSocketType = Stream}
      getAddrInfo (Just hints) (Just host) (Just port)

    tryOpen :: Exception e => Either e Handle -> AddrInfo -> IO (Either e Handle)
    tryOpen h@(Right _) _ = return h
    tryOpen (Left _) addr = try $ open addr

    open :: AddrInfo -> IO Handle
    open addr = do
      sock <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
      connect sock $ addrAddress addr
      getSocketHandle sock

Server.hs

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Network.Socket
import System.IO
import Util

main :: IO ()
main = do
  putStrLn "listening to 5000..."
  runTCPServer pong

pong :: Handle -> IO ()
pong h = forever $ do
  line <- getLn h
  print line
  putLn h "pong"

runTCPServer :: (Handle -> IO ()) -> IO ()
runTCPServer server =
  bracket startTCPServer close $ \sock -> forever $ do
    h <- acceptTCPConn sock
    putStrLn "client connected"
    forkFinally (server h) $ \_ -> do
      hClose h
      putStrLn "client disconnected"

startTCPServer :: IO Socket
startTCPServer = withSocketsDo $ resolve >>= open
  where
    resolve = do
      let hints = defaultHints {addrFlags = [AI_PASSIVE], addrSocketType = Stream}
      head <$> getAddrInfo (Just hints) Nothing (Just "5000")
    open addr = do
      sock <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
      setSocketOption sock ReuseAddr 1
      withFdSocket sock setCloseOnExecIfNeeded
      bind sock $ addrAddress addr
      listen sock 1024
      return sock

acceptTCPConn :: Socket -> IO Handle
acceptTCPConn sock = do
  (conn, _) <- accept sock
  getSocketHandle conn

Util.hs

{-# LANGUAGE OverloadedStrings #-}

module Util where

import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import Network.Socket
import System.IO

getSocketHandle :: Socket -> IO Handle
getSocketHandle conn = do
  h <- socketToHandle conn ReadWriteMode
  hSetBinaryMode h True
  hSetNewlineMode h NewlineMode {inputNL = CRLF, outputNL = CRLF}
  hSetBuffering h LineBuffering
  return h

putLn :: Handle -> ByteString -> IO ()
putLn h = B.hPut h . (<> "\r\n")

getLn :: Handle -> IO ByteString
getLn h = trim_cr <$> B.hGetLine h
  where
    trim_cr "" = ""
    trim_cr s = if B.last s == '\r' then B.init s else s