A stub for TCP/IP handling that I did with the perspective of doing something with Pixelflut for ~1h back in 2016

2016-09-04.Network.hs

module Network(
  defaultPort, Connection(..), send, recv, close,
  Listening(..), stopListening, serve, connect,
  serve', connect'
)where
{-|
  This module handles opening/accepting connections on/to an interface/ip,
  to send/receive Data.Binary
|-}
import Control.Monad
import Data.Binary (Binary(..), Word8)
import Data.ByteString.Lazy (ByteString)
import Network.Socket (HostName, PortNumber)
import System.IO (Handle)
import qualified Control.Concurrent as Concurrent
import qualified Data.Binary as Binary
import qualified Data.ByteString.Lazy as Lazy
import qualified Data.Maybe as Maybe
import qualified Network.Simple.TCP as TCP
import qualified Network.Socket as Socket
import qualified System.IO as IO

-- Largest prime below 65536
defaultPort = 65521 :: PortNumber

-- Connection relevant data:
data Connection = Connection {
    socket :: Socket.Socket
  , handle :: Handle
  , addr   :: Socket.SockAddr
  }

mkConnection :: (Socket.Socket, Socket.SockAddr) -> IO Connection
mkConnection (s,a) = do
  h <- Socket.socketToHandle s IO.ReadWriteMode
  return $ Connection s h a

send :: Binary b => Connection -> b -> IO ()
send c b = do
  let bytes = Binary.encode b
      send' = Lazy.hPut (handle c)
  send' . Binary.encode $ sizeOf bytes
  send' bytes

recv :: Binary b => Connection -> IO b
recv c = do
  let bCount = sizeOf $ Binary.encode (23 :: Int)
      recv'  = Lazy.hGet (handle c)
  size <- Binary.decode <$> recv' bCount
  Binary.decode <$> recv' size

-- Helper for {send,recv}
sizeOf :: ByteString -> Int
sizeOf = fromIntegral . Lazy.length

close :: Connection -> IO ()
close = IO.hClose . handle

-- Encapsulate ThreadIds that listen on different HostNames
newtype Listening = CListen {fromListen :: [Concurrent.ThreadId]}

stopListening :: Listening -> IO ()
stopListening = mapM_ Concurrent.killThread . fromListen

serve :: (Connection -> IO ()) -> [HostName] -> IO Listening
serve f hs = serve' f . zip hs $ repeat defaultPort

serve' :: (Connection -> IO ()) -> [(HostName, PortNumber)] -> IO Listening
serve' f = fmap CListen . go (f <=< mkConnection)
  where
    go :: ((Socket.Socket, Socket.SockAddr) -> IO ()) -> [(HostName, PortNumber)] -> IO [Concurrent.ThreadId]
    go f = mapM $ \(hn, p) -> Concurrent.forkIO $ TCP.serve (TCP.Host hn) (show p) f

connect :: HostName -> (Connection -> IO a) -> IO a
connect = connect' `flip` defaultPort

connect' :: HostName -> PortNumber -> (Connection -> IO a) -> IO a
connect' hn p f = TCP.connect hn (show p) (f <=< mkConnection)