Udp.hs

#!/usr/bin/env stack
-- stack --resolver lts-16.4 script --package network --package bytestring --package split

-- A minimalistic UDP library

{-# LANGUAGE OverloadedStrings, RecordWildCards #-}
module Udp ( withServer
           , parseAddr
           , Send
           , Recv
           , HostPort
           , main
           ) where

import           Control.Concurrent        (forkIO, threadDelay)
import           Control.Exception         (bracket)
import           Control.Monad             (forever)
import           Data.ByteString           (ByteString)
import qualified Data.ByteString.Char8     as BS
import           Data.List.Split           (splitOn)
import           Network.Socket            (AddrInfo (..), Family (..),
                                            SockAddr (..), SocketType (..),
                                            bind, close, defaultHints,
                                            getAddrInfo, getSocketName, socket,
                                            withSocketsDo)
import           Network.Socket.ByteString (recvFrom, sendAllTo)
import           System.Environment        (getArgs)
import           System.IO                 (BufferMode (LineBuffering),
                                            hSetBuffering, stdout)

-- API
type Send     = SockAddr -> ByteString -> IO ()
type Recv     = IO (ByteString, SockAddr)
type HostPort = String -- 1.2.3.4:567

parseAddr :: HostPort -> IO SockAddr
parseAddr s = addrAddress <$> parseInfo s

withServer :: String -> HostPort -> (SockAddr -> Send -> Recv -> IO a) -> IO a
withServer side hostPort f = do
    AddrInfo{..} <- parseInfo hostPort
    bracket
        (socket addrFamily addrSocketType addrProtocol)
        close
        $ \s -> do
            bind s addrAddress
            boundAddr <- getSocketName s

            putStrLn $ side ++ " ** Serving on " ++ show boundAddr
            f boundAddr (send s) (recv s)
  where
    send s to bs = do
        putStrLn $ side ++ " >> " ++ show (BS.unpack bs) ++ " to   " ++ show to
        sendAllTo s bs to
    recv s = do
        sfrom@(bs, from) <- recvFrom s 4096
        putStrLn $ side ++ " << " ++ show (BS.unpack bs) ++ " from " ++ show from
        return sfrom

parseInfo :: HostPort -> IO AddrInfo
parseInfo hostPort = do
    let hints = defaultHints
            { addrFamily     = AF_INET
            , addrSocketType = Datagram }

        (mHost, mPort) = case splitOn ":" hostPort of
            [host]       -> (Just host, Nothing)
            [host, port] -> (Just host, Just port)
            _            -> error $ "Couldn't parse HostPort: " ++ hostPort

    -- TODO check if empty
    head <$> getAddrInfo (Just hints) mHost mPort

-- Example
--
-- ./Udp.hs
-- A ** Serving on 0.0.0.0:7000
-- B ** Serving on 0.0.0.0:52498
-- B >> "Hello, world!" to   127.0.0.1:7000
-- A << "Hello, world!" from 127.0.0.1:52498
-- A >> "Hello, world!" to   127.0.0.1:52498
-- B << "Hello, world!" from 127.0.0.1:7000
--
main :: IO ()
main = withSocketsDo $ do
    hSetBuffering stdout LineBuffering

    forkIO $ withServer "A" "0.0.0.0:7000"  $ \_self send recv ->
        forever $ do
            (bs, from) <- recv
            send from bs

    threadDelay 1000 -- 1ms

    withServer "B" "0.0.0.0" $ \_self send recv -> do
        forkIO $ forever $ do
            (_bs, _from) <- recv
            return ()

        a <- parseAddr "127.0.0.1:7000"

        forever $ do
            send a "Hello, world!"
            threadDelay 1000000 -- 1s