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

59
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,36 +85,41 @@ 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.close (N.bindPortUDP (fromIntegral port) (fromString host))
(runEventLoop notebook) N.close
(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 let appData = UdpAppData { appAddr = addr
clientCtx <- H.lookup addr <$> readMVar clientMapM , 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 -> (HostName, PortNumber) -> (HostName, PortNumber) -> (WS.Connection -> IO ()) -> IO ()
runTunnelingClient proto (wsHost, wsPort) (remoteHost, remotePort) app = do runTunnelingClient proto (wsHost, wsPort) (remoteHost, remotePort) app = do