-- | Simple functions to run UDP clients and servers.
module Network.Run.UDP (
    runUDPClient,
    runUDPServer,
    runUDPServerFork,
) where

import Control.Concurrent (forkFinally, forkIO)
import qualified Control.Exception as E
import Control.Monad (forever, void)
import Data.ByteString (ByteString)
import qualified Data.List.NonEmpty as NE
import Network.Socket
import Network.Socket.ByteString

import Network.Run.Core

-- | Running a UDP client with a socket.
--   The client action takes a socket and
--   server's socket address.
--   They should be used with 'sendTo'.
runUDPClient :: HostName -> ServiceName -> (Socket -> SockAddr -> IO a) -> IO a
runUDPClient :: forall a.
HostName -> HostName -> (Socket -> SockAddr -> IO a) -> IO a
runUDPClient HostName
host HostName
port Socket -> SockAddr -> IO a
client = IO a -> IO a
forall a. IO a -> IO a
withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
    AddrInfo
addr <- SocketType
-> Maybe HostName -> HostName -> [AddrInfoFlag] -> IO AddrInfo
resolve SocketType
Datagram (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
host) HostName
port [AddrInfoFlag
AI_ADDRCONFIG]
    let sockAddr :: SockAddr
sockAddr = AddrInfo -> SockAddr
addrAddress AddrInfo
addr
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (AddrInfo -> IO Socket
openSocket AddrInfo
addr) Socket -> IO ()
close ((Socket -> IO a) -> IO a) -> (Socket -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> Socket -> SockAddr -> IO a
client Socket
sock SockAddr
sockAddr

-- | Running a UDP server with an open socket in a single Haskell thread.
runUDPServer :: Maybe HostName -> ServiceName -> (Socket -> IO a) -> IO a
runUDPServer :: forall a. Maybe HostName -> HostName -> (Socket -> IO a) -> IO a
runUDPServer Maybe HostName
mhost HostName
port Socket -> IO a
server = IO a -> IO a
forall a. IO a -> IO a
withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
    AddrInfo
addr <- SocketType
-> Maybe HostName -> HostName -> [AddrInfoFlag] -> IO AddrInfo
resolve SocketType
Datagram Maybe HostName
mhost HostName
port [AddrInfoFlag
AI_PASSIVE]
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (AddrInfo -> IO Socket
openServerSocket AddrInfo
addr) Socket -> IO ()
close Socket -> IO a
server

-- | Running a UDP server with a connected socket in each Haskell thread.
--   The first request is given to the server.
--   Suppose that the server is serving on __addrS:portS__ and
--   a client connects to the service from __addrC:portC__.
--   A connected socket is created by binding to __*:portS__ and
--   connecting to __addrC:portC__,
--   resulting in __(UDP,addrS:portS,addrC:portC)__ where
--   __addrS__ is given magically.
--   This approach is fragile due to NAT rebidings.
runUDPServerFork
    :: [HostName] -> ServiceName -> (Socket -> ByteString -> IO ()) -> IO ()
runUDPServerFork :: [HostName] -> HostName -> (Socket -> ByteString -> IO ()) -> IO ()
runUDPServerFork [] HostName
_ Socket -> ByteString -> IO ()
_ = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
runUDPServerFork (HostName
h : [HostName]
hs) HostName
port Socket -> ByteString -> IO ()
server = do
    (HostName -> IO ThreadId) -> [HostName] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId)
-> (HostName -> IO ()) -> HostName -> IO ThreadId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HostName -> IO ()
forall {b}. HostName -> IO b
run) [HostName]
hs
    HostName -> IO ()
forall {b}. HostName -> IO b
run HostName
h
  where
    run :: HostName -> IO b
run HostName
host = do
        HostName -> IO ()
labelMe (HostName -> IO ()) -> HostName -> IO ()
forall a b. (a -> b) -> a -> b
$ HostName
"UDP server for " HostName -> HostName -> HostName
forall a. [a] -> [a] -> [a]
++ HostName
h
        Maybe HostName -> HostName -> (Socket -> IO b) -> IO b
forall a. Maybe HostName -> HostName -> (Socket -> IO a) -> IO a
runUDPServer (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
host) HostName
port ((Socket -> IO b) -> IO b) -> (Socket -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Socket
lsock -> IO () -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO b) -> IO () -> IO b
forall a b. (a -> b) -> a -> b
$ do
            (ByteString
bs0, SockAddr
peeraddr) <- Socket -> Int -> IO (ByteString, SockAddr)
recvFrom Socket
lsock Int
2048
            let family :: Family
family = case SockAddr
peeraddr of
                    SockAddrInet{} -> Family
AF_INET
                    SockAddrInet6{} -> Family
AF_INET6
                    SockAddr
_ -> HostName -> Family
forall a. HasCallStack => HostName -> a
error HostName
"family"
                hints :: AddrInfo
hints =
                    AddrInfo
defaultHints
                        { addrSocketType = Datagram
                        , addrFamily = family
                        , addrFlags = [AI_PASSIVE]
                        }
            AddrInfo
addr <- NonEmpty AddrInfo -> AddrInfo
forall a. NonEmpty a -> a
NE.head (NonEmpty AddrInfo -> AddrInfo)
-> IO (NonEmpty AddrInfo) -> IO AddrInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (NonEmpty AddrInfo)
forall (t :: * -> *).
GetAddrInfo t =>
Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (t AddrInfo)
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) Maybe HostName
forall a. Maybe a
Nothing (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
port)
            Socket
s <- AddrInfo -> IO Socket
openServerSocket AddrInfo
addr
            Socket -> SockAddr -> IO ()
connect Socket
s SockAddr
peeraddr
            IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> (Either SomeException () -> IO ()) -> IO ThreadId
forall a. IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinally (HostName -> IO ()
labelMe HostName
"UDP server" IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Socket -> ByteString -> IO ()
server Socket
s ByteString
bs0) (\Either SomeException ()
_ -> Socket -> IO ()
close Socket
s)