asakul
/
qs-tunnel
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

Multiple upstreams

master
Denis Tereshkin 6 years ago
parent
4bbbd6e19f
commit
e6338fe7db
2 changed files with 115 additions and 71 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 4
      qs-tunnel.cabal
  2. 178
      src/Main.hs

4
qs-tunnel.cabal
Unescape View File

@ -1,5 +1,5 @@
name: qs-tunnel name: qs-tunnel
version: 0.1.0.0 version: 0.2.0.0
synopsis: Quotesource tunnel proxy synopsis: Quotesource tunnel proxy
-- description: -- description:
homepage: https://github.com/asakul/qs-tunnel#readme homepage: https://github.com/asakul/qs-tunnel#readme
@ -7,7 +7,7 @@ license: BSD3
license-file: LICENSE license-file: LICENSE
author: Denis Tereshkin author: Denis Tereshkin
maintainer: denis@kasan.ws maintainer: denis@kasan.ws
copyright: 2017 Denis Tereshkin copyright: 2017-2019 Denis Tereshkin
category: Web category: Web
build-type: Simple build-type: Simple
cabal-version: >=1.10 cabal-version: >=1.10

178
src/Main.hs
Unescape View File

@ -2,46 +2,63 @@
module Main where module Main where
import qualified Data.Text as T import Data.Aeson
import qualified Data.ByteString as B import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8 import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Lazy as BL import qualified Data.ByteString.Lazy as BL
import qualified Data.List as L import Data.IORef
import Data.List.NonEmpty import qualified Data.List as L
import Data.IORef import Data.List.NonEmpty
import Data.Time.Clock import qualified Data.Text as T
import Data.Aeson import Data.Time.Clock
import ATrade.QuoteSource.Server import ATrade.QuoteSource.Client
import ATrade.QuoteSource.Client import ATrade.QuoteSource.Server
import Control.Monad import Control.Concurrent
import Control.Monad.Loops import Control.Monad
import Control.Monad.Loops
import System.IO
import System.Log.Logger import System.IO
import System.Log.Handler.Simple import System.Log.Formatter
import System.Log.Handler (setFormatter) import System.Log.Handler (setFormatter)
import System.Log.Formatter import System.Log.Handler.Simple
import System.ZMQ4 import System.Log.Logger
import System.ZMQ4.ZAP import System.ZMQ4
import System.ZMQ4.ZAP
data UpstreamConfig = UpstreamConfig
{
ucEndpoint :: T.Text,
ucCertificatePath :: Maybe FilePath
} deriving (Show, Eq)
instance FromJSON UpstreamConfig where
parseJSON (Object o) =
UpstreamConfig <$>
o .: "endpoint" <*>
o .:? "certificate"
data Config = Config data Config = Config
{ {
confDownstreamEp :: T.Text, confDownstreamEp :: T.Text,
confWhitelistIps :: [T.Text], confDownstreamCertificatePath :: Maybe FilePath,
confBlacklistIps :: [T.Text], confClientCertificates :: [FilePath],
confUpstreamEp :: T.Text, confWhitelistIps :: [T.Text],
confTimeout :: Integer confBlacklistIps :: [T.Text],
confUpstreams :: [UpstreamConfig],
confTimeout :: Integer
} deriving (Show, Eq) } deriving (Show, Eq)
instance FromJSON Config where instance FromJSON Config where
parseJSON (Object o) = parseJSON (Object o) =
Config <$> Config <$>
o .: "downstream" <*> o .: "downstream" <*>
o .:? "whitelist" .!= [] <*> o .:? "downstream_certificate" <*>
o .:? "blacklist" .!= [] <*> o .: "client_certificates" <*>
o .: "upstream" <*> o .:? "whitelist" .!= [] <*>
o .:? "blacklist" .!= [] <*>
o .: "upstreams" <*>
o .: "timeout" o .: "timeout"
parseJSON _ = fail "Expected object" parseJSON _ = fail "Expected object"
@ -62,43 +79,70 @@ main = do
eConf <- eitherDecode . BL.fromStrict <$> B.readFile "qs-tunnel.conf" eConf <- eitherDecode . BL.fromStrict <$> B.readFile "qs-tunnel.conf"
case eConf of case eConf of
Left errMsg -> error errMsg Left errMsg -> error errMsg
Right conf -> runWithConfig conf Right conf -> runWithConfig conf
runWithConfig conf = do runWithConfig conf = do
withContext $ \ctx -> withContext $ \ctx ->
withSocket ctx Pub $ \downstream -> do withZapHandler ctx $ \zap -> do
bind downstream $ T.unpack $ confDownstreamEp conf withSocket ctx Pub $ \downstream -> do
setZapDomain (restrict "global") downstream
forever $ withSocket ctx Sub $ \upstream -> do zapSetBlacklist zap "global" $ confBlacklistIps conf
infoM "main" $ "Connecting to: " ++ (T.unpack $ confUpstreamEp conf) zapSetWhitelist zap "global" $ confWhitelistIps conf
connect upstream $ T.unpack $ confUpstreamEp conf bind downstream $ T.unpack $ confDownstreamEp conf
subscribe upstream B.empty case (confDownstreamCertificatePath conf) of
now <- getCurrentTime Just certPath -> do
lastHeartbeat <- newIORef now eCert <- loadCertificateFromFile certPath
lastHeartbeatSent <- newIORef now case eCert of
infoM "main" "Starting proxy loop" Left err -> errorM "main" $ "Unable to load certificate: " ++ certPath
whileM (notTimeout lastHeartbeat conf) $ do Right cert -> do
evs <- poll 200 [Sock upstream [In] Nothing] zapSetServerCertificate cert downstream
sendHeartbeatIfNeeded lastHeartbeatSent downstream forM_ (confClientCertificates conf) (addCertificate zap)
unless (null (L.head evs)) $ do _ -> return ()
incoming <- receiveMulti upstream
case incoming of forM_ (confUpstreams conf) $ \upstreamConf -> forkIO $ do
x:xs -> do forever $ withSocket ctx Sub $ \upstream -> do
now <- getCurrentTime infoM "main" $ "Connecting to: " ++ (T.unpack $ ucEndpoint upstreamConf)
writeIORef lastHeartbeat now case (ucCertificatePath upstreamConf) of
sendMulti downstream $ x :| xs Just certPath -> do
eCert <- loadCertificateFromFile certPath
case eCert of
Left err -> errorM "main" $ "Unable to load certificate: " ++ certPath
Right cert -> zapApplyCertificate cert upstream
_ -> return () _ -> return ()
where connect upstream $ T.unpack $ ucEndpoint upstreamConf
notTimeout ref conf = do subscribe upstream B.empty
now <- getCurrentTime now <- getCurrentTime
lastHb <- readIORef ref lastHeartbeat <- newIORef now
return $ now `diffUTCTime` lastHb < (fromInteger . confTimeout) conf lastHeartbeatSent <- newIORef now
infoM "main" "Starting proxy loop"
sendHeartbeatIfNeeded lastHbSent sock = do whileM (notTimeout lastHeartbeat conf) $ do
now <- getCurrentTime evs <- poll 200 [Sock upstream [In] Nothing]
last <- readIORef lastHbSent sendHeartbeatIfNeeded lastHeartbeatSent downstream
when (now `diffUTCTime` last > 1) $ do unless (null (L.head evs)) $ do
send sock [] $ B8.pack "SYSTEM#HEARTBEAT" incoming <- receiveMulti upstream
writeIORef lastHbSent now case incoming of
x:xs -> do
now <- getCurrentTime
writeIORef lastHeartbeat now
sendMulti downstream $ x :| xs
_ -> return ()
where
notTimeout ref conf = do
now <- getCurrentTime
lastHb <- readIORef ref
return $ now `diffUTCTime` lastHb < (fromInteger . confTimeout) conf
sendHeartbeatIfNeeded lastHbSent sock = do
now <- getCurrentTime
last <- readIORef lastHbSent
when (now `diffUTCTime` last > 1) $ do
send sock [] $ B8.pack "SYSTEM#HEARTBEAT"
writeIORef lastHbSent now
addCertificate zap clientCertPath = do
eClientCert <- loadCertificateFromFile clientCertPath
case eClientCert of
Left err -> errorM "main" $ "Unable to load client certificate: " ++ clientCertPath
Right clientCert -> zapAddClientCertificate zap "global" clientCert

Write Preview
Loading…
Cancel
Save