diff --git a/src/Lib.hs b/src/Lib.hs index 4f93272..b68fc0b 100644 --- a/src/Lib.hs +++ b/src/Lib.hs @@ -1,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE NoImplicitPrelude #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -27,7 +28,6 @@ import qualified Network.Socket.ByteString as N import qualified Network.WebSockets as WS import qualified Network.WebSockets.Stream as WS -import Data.Maybe import Network.Connection (Connection, ConnectionParams (..), TLSSettings (..), connectTo, connectionGetChunk, connectionPut, @@ -85,36 +85,41 @@ runUDPClient (host, port) app = do runUDPServer :: (HostName, PortNumber) -> (UdpAppData -> IO ()) -> IO () runUDPServer (host, port) app = do putStrLn $ "WAIT for datagrames on " <> tshow host <> ":" <> tshow port - notebook <- newMVar mempty - bracket (N.bindPortUDP (fromIntegral port) (fromString host)) - N.close - (runEventLoop notebook) + clientsCtx <- newMVar mempty + void $ bracket + (N.bindPortUDP (fromIntegral port) (fromString host)) + N.close + (runEventLoop clientsCtx) putStrLn "CLOSE tunnel" where - runEventLoop :: MVar (H.HashMap N.SockAddr UdpAppData) -> N.Socket -> IO () - runEventLoop clientMapM socket = forever $ do - (payload, addr) <- N.recvFrom socket 4096 - clientCtx <- H.lookup addr <$> readMVar clientMapM + addNewClient clientsCtx socket addr payload = do + sem <- newMVar payload + let appData = UdpAppData { appAddr = addr + , appSem = sem + , appRead = takeMVar sem + , appWrite = \payload' -> void $ N.sendTo socket payload' addr + } + void $ withMVar clientsCtx (return . H.insert addr appData) + return appData + + removeClient clientsCtx clientCtx = do + void $ withMVar clientsCtx (return . H.delete (appAddr clientCtx)) + putStrLn "TIMEOUT connection" + + pushDataToClient clientCtx = putMVar (appSem clientCtx) + + runEventLoop clientsCtx socket = forever $ do + (payload, addr) <- N.recvFrom socket 4096 + clientCtx <- H.lookup addr <$> readMVar clientsCtx + + case clientCtx of + Just clientCtx' -> pushDataToClient clientCtx' payload + _ -> void $ async $ bracket + (addNewClient clientsCtx socket addr payload) + (removeClient clientsCtx) + (timeout (30 * 10^(6 :: Int)) . app) - if isJust clientCtx - then putMVar (appSem $ fromJust clientCtx) payload - else - void $ async $ bracket - (do sem <- newMVar payload - let appData = UdpAppData { appAddr = addr - , appSem = sem - , appRead = takeMVar sem - , appWrite = \payload' -> void $ N.sendTo socket payload' addr - } - void $ withMVar clientMapM (return . H.insert addr appData) - return appData - ) - (\appData' -> do - void $ withMVar clientMapM (return . H.delete (appAddr appData')) - putStrLn "TIMEOUT connection" - ) - (timeout (30 * 10^(6 :: Int)) . app) runTunnelingClient :: Proto -> (HostName, PortNumber) -> (HostName, PortNumber) -> (WS.Connection -> IO ()) -> IO () runTunnelingClient proto (wsHost, wsPort) (remoteHost, remotePort) app = do