hex/wstunnel
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Improve readability

Browse source
This commit is contained in:
Erèbe 2016-05-16 22:19:34 +02:00
parent 499f872d1a
commit 58a313f354
1 changed files with 32 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
src/Lib.hs
View file

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