1 files changed, 81 insertions, 0 deletions
diff --git a/fig-utils/src/Fig/Utils/Net.hs b/fig-utils/src/Fig/Utils/Net.hs
new file mode 100644
index 0000000..d0b1890
--- /dev/null
+++ b/fig-utils/src/Fig/Utils/Net.hs
@@ -0,0 +1,81 @@
+module Fig.Utils.Net
+  ( server
+  , client
+  ) where
+
+import Fig.Prelude
+
+import System.IO (IOMode(..), BufferMode(..), hClose, hSetBuffering)
+
+import qualified Control.Concurrent as Conc
+
+import qualified Network.Socket as Sock
+
+newtype FigNetException = FigNetException Text
+  deriving (Show, Eq, Ord)
+instance Exception FigNetException
+
+resolveAddr :: forall m.
+  (MonadIO m, MonadThrow m) =>
+  (Maybe Text, Text) ->
+  Bool ->
+  m Sock.AddrInfo
+resolveAddr (host, port) serv = do
+  maddr <- liftIO $ headMay <$> Sock.getAddrInfo
+    (Just $ Sock.defaultHints
+     { Sock.addrFlags = [Sock.AI_PASSIVE | serv]
+     , Sock.addrSocketType = Sock.Stream
+     }
+    )
+    (unpack <$> host)
+    (Just $ unpack port)
+  maybe (throwM $ FigNetException "Failed to resolve address") pure maddr
+
+server :: forall m.
+  (MonadIO m, MonadThrow m, MonadMask m) =>
+  (Maybe Text, Text) ->
+  m (Handle -> Sock.SockAddr -> (IO (), IO ())) ->
+  m ()
+server loc onConn = do
+  addr <- resolveAddr loc True
+  bracket (liftIO $ Sock.openSocket addr) (liftIO . Sock.close) \sock -> do
+    liftIO $ Sock.setSocketOption sock Sock.ReuseAddr 1
+    liftIO $ Sock.withFdSocket sock Sock.setCloseOnExecIfNeeded
+    liftIO $ Sock.bind sock $ Sock.addrAddress addr
+    liftIO $ Sock.listen sock 4096
+    log $ "Listening on " <> tshow (Sock.addrAddress addr)
+    forever do
+      let toHandle = bracketOnError (liftIO $ Sock.accept sock) (liftIO . Sock.close . fst) \(conn, peer) ->
+            liftIO $ (,peer) <$> Sock.socketToHandle conn ReadWriteMode
+      bracketOnError toHandle (liftIO . hClose . fst) \(hdl, peer) -> do
+        liftIO $ log $ "Client " <> tshow peer <> " connected"
+        liftIO $ hSetBuffering hdl LineBuffering
+        (handler, cleanup) <- ($ peer) . ($ hdl) <$> onConn
+        liftIO $ Conc.forkFinally handler \res -> do
+          case res of
+            Right _ -> log $ "Client " <> tshow peer <> " disconnected"
+            Left err -> log $ "Client " <> tshow peer <> " disconnected: " <> tshow err
+          cleanup
+          hClose hdl
+
+client :: forall m.
+  (MonadIO m, MonadThrow m, MonadMask m) =>
+  (Text, Text) ->
+  m (Handle -> (m (), m ())) ->
+  m ()
+client loc onConn = do
+  addr <- resolveAddr (first Just loc) False
+  let openConnectHandle = do
+        bracketOnError (liftIO $ Sock.openSocket addr) (liftIO . Sock.close) \sock -> do
+          liftIO . Sock.connect sock $ Sock.addrAddress addr
+          hdl <- liftIO $ Sock.socketToHandle sock ReadWriteMode
+          (handler, cleanup) <- ($ hdl) <$> onConn
+          pure (hdl, handler, cleanup)
+  bracket openConnectHandle
+    ( \(hdl, _, cleanup) -> do
+        cleanup 
+        liftIO $ hClose hdl
+    )
+    ( \(_, handler, _) -> do
+        handler
+    )