natskell-1.0.0.1: test/Unit/ConnectionSpec.hs
{-# LANGUAGE OverloadedStrings #-}
module ConnectionSpec (spec) where
import Auth.None (auth)
import qualified Client
import Control.Concurrent
import Control.Concurrent.STM
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.IORef (newIORef, readIORef, writeIORef)
import Handshake.Nats
( HandshakeError (HandshakeAuthError, HandshakeTLSError, HandshakeTimeout)
, performHandshake
)
import Lib.Logger
( LogLevel (Debug)
, LoggerConfig (LoggerConfig)
, newLogContext
)
import Network.Connection (connectionApi)
import Network.Connection.Core (Transport (..), pointTransport)
import Network.ConnectionAPI
( closeReader
, newConn
, readData
, reader
)
import Parser.API
( ParseStep (DropPrefix, Emit, NeedMore, Reject)
, ParsedMessage (ParsedInfo, ParsedPing, ParsedPong)
, ParserAPI (ParserAPI)
)
import qualified Parser.Attoparsec as Attoparsec
import qualified Queue.API as Queue
import Queue.TransactionalQueue (newQueue)
import Router.Nats
( RouteDirective (RouteContinue)
, routeMessage
)
import State.Store
( newClientState
, pushPingAction
, queue
, readServerInfo
)
import State.Types (ClientConfig (..))
import Subscription.Store (newSubscriptionStore)
import Subscription.Types (defaultPendingLimits)
import System.Timeout (timeout)
import Test.Hspec
import Transformers.Transformers (Transformer (transform))
import qualified Types.Connect as Connect
import Types.Info (Info (Info))
import Types.Ping (Ping (Ping))
import Types.Pong (Pong (Pong))
import Types.TLS (TLSConfig (..), defaultTLSConfig)
spec :: Spec
spec = do
describe "TLS configuration" $ do
it "does not expose client private keys when rendered" $ do
let config = defaultTLSConfig
{ tlsClientCertificate = Just ("certificate", "private-key")
, tlsRootCertificates = ["private-root"]
}
show config `shouldNotContain` "private-key"
show config `shouldNotContain` "private-root"
show config `shouldContain` "tlsRootCertificates = 1 configured"
describe "Connection reader" $ do
it "unblocks a blocking read when closeReader is called" $ do
conn <- newConn connectionApi
started <- newEmptyMVar
blocker <- (newEmptyMVar :: IO (MVar BS.ByteString))
let transport = Transport
{ transportRead = \_ -> putMVar started () >> takeMVar blocker
, transportWrite = \_ -> pure ()
, transportWriteLazy = \_ -> pure ()
, transportFlush = pure ()
, transportClose = pure ()
, transportUpgrade = Nothing
}
pointTransport conn transport
resultVar <- newEmptyMVar
_ <- forkIO $ readData (reader connectionApi) conn 1 >>= putMVar resultVar
_ <- takeMVar started
closeReader (reader connectionApi) conn
result <- timeout 1000000 (takeMVar resultVar)
result `shouldBe` Just (Left "Read operation is blocked")
describe "Router ping lifecycle" $ do
it "responds to server PING by enqueueing PONG" $ do
state <- newTestState
store <- newSubscriptionStore defaultPendingLimits (pure ())
directive <- routeMessage state store (ParsedPing Ping)
directive `shouldBe` RouteContinue
queued <- Queue.dequeue (queue state)
case queued of
Right item ->
LBS.toStrict (transform item) `shouldBe` "PONG\r\n"
Left err ->
expectationFailure ("expected queued PONG, got queue error: " ++ err)
it "runs exactly one pending ping action for each PONG" $ do
state <- newTestState
store <- newSubscriptionStore defaultPendingLimits (pure ())
firstActionRan <- newIORef False
secondActionRan <- newIORef False
pushPingAction state (writeIORef firstActionRan True)
pushPingAction state (writeIORef secondActionRan True)
firstDirective <- routeMessage state store (ParsedPong Pong)
firstDirective `shouldBe` RouteContinue
readIORef firstActionRan `shouldReturn` True
readIORef secondActionRan `shouldReturn` False
secondDirective <- routeMessage state store (ParsedPong Pong)
secondDirective `shouldBe` RouteContinue
readIORef secondActionRan `shouldReturn` True
describe "Handshake" $ do
it "returns a typed error when no servers are configured" $ do
result <- Client.newClient [] [Client.withConnectionAttempts 1]
case result of
Left Client.ConnectNoServers -> pure ()
Left err -> expectationFailure ("unexpected connection error: " ++ show err)
Right _ -> expectationFailure "client connected without a server"
it "accepts an incremental parser backend for the initial INFO frame" $ do
state <- newTestState
conn <- newConn connectionApi
writes <- newTVarIO []
transport <-
newScriptedTransport
[ "INF"
, "O {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n"
, "PONG\r\n"
]
writes
pointTransport conn transport
result <- performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1"
result `shouldBe` Right ()
readServerInfo state `shouldReturn` Just testInfo
readTVarIO writes `shouldReturn`
[LBS.toStrict (transform testConnect <> transform Ping)]
it "accepts a parser backend that drops an invalid prefix before INFO" $ do
state <- newTestState
conn <- newConn connectionApi
writes <- newTVarIO []
transport <-
newScriptedTransport
[ "XIN"
, "FO {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n"
, "PONG\r\n"
]
writes
pointTransport conn transport
result <- performHandshake connectionApi dropPrefixInfoParser state auth conn "127.0.0.1"
result `shouldBe` Right ()
readServerInfo state `shouldReturn` Just testInfo
readTVarIO writes `shouldReturn`
[LBS.toStrict (transform testConnect <> transform Ping)]
it "does not report readiness until the server sends PONG" $ do
state <- newTestStateWithTimeout 20000
conn <- newConn connectionApi
writes <- newTVarIO []
blocker <- newEmptyMVar
let transport = Transport
{ transportRead = \_ -> do
first <- atomically $ do
written <- readTVar writes
pure (null written)
if first then pure infoFrame else takeMVar blocker
, transportWrite = \bytes ->
atomically $ modifyTVar' writes (<> [bytes])
, transportWriteLazy = \bytes ->
atomically $ modifyTVar' writes (<> [LBS.toStrict bytes])
, transportFlush = pure ()
, transportClose = pure ()
, transportUpgrade = Nothing
}
pointTransport conn transport
result <- performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1"
result `shouldBe` Left HandshakeTimeout
it "accepts +OK before the handshake PONG" $ do
state <- newTestState
conn <- newConn connectionApi
writes <- newTVarIO []
transport <- newScriptedTransport [infoFrame, "+OK\r\nPONG\r\n"] writes
pointTransport conn transport
result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1"
result `shouldBe` Right ()
it "returns authentication errors received before PONG" $ do
state <- newTestState
conn <- newConn connectionApi
writes <- newTVarIO []
transport <- newScriptedTransport
[infoFrame, "-ERR 'Authorization Violation'\r\n"] writes
pointTransport conn transport
result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1"
case result of
Left (HandshakeAuthError _) -> pure ()
other -> expectationFailure ("expected auth error, got: " ++ show other)
it "rejects requested TLS when the server does not advertise it" $ do
state <- newTestStateWithConfig $ \cfg ->
cfg { tlsConfig = Just defaultTLSConfig }
conn <- newConn connectionApi
writes <- newTVarIO []
transport <- newScriptedTransport [infoFrame] writes
pointTransport conn transport
result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1"
case result of
Left (HandshakeTLSError _) -> pure ()
other -> expectationFailure ("expected TLS error, got: " ++ show other)
newTestState = do
newTestStateWithConfig id
newTestStateWithTimeout timeoutMicros = do
newTestStateWithConfig $ \cfg -> cfg { connectTimeoutMicros = timeoutMicros }
newTestStateWithConfig updateConfig = do
queue <- newQueue
ctx <- newLogContext
conn <- newConn connectionApi
logger <- newSilentLogger
newClientState (updateConfig (testConfig logger)) queue conn ctx
newSilentLogger :: IO LoggerConfig
newSilentLogger = do
lock <- newTMVarIO ()
pure (LoggerConfig Debug (\_ -> pure ()) lock)
testConfig :: LoggerConfig -> ClientConfig
testConfig logger =
ClientConfig
{ connectionAttempts = 1
, connectTimeoutMicros = 1000000
, callbackConcurrency = 1
, messageLimit = 1024 * 1024
, connectConfig = testConnect
, loggerConfig = logger
, tlsConfig = Nothing
, exitAction = const (pure ())
, connectOptions = []
}
testConnect :: Connect.Connect
testConnect =
Connect.Connect
{ Connect.verbose = False
, Connect.pedantic = True
, Connect.tls_required = False
, Connect.auth_token = Nothing
, Connect.user = Nothing
, Connect.pass = Nothing
, Connect.name = Nothing
, Connect.lang = "haskell"
, Connect.version = "0.1.0"
, Connect.protocol = Nothing
, Connect.echo = Just True
, Connect.sig = Nothing
, Connect.jwt = Nothing
, Connect.nkey = Nothing
, Connect.no_responders = Just True
, Connect.headers = Just True
}
testInfo :: Info
testInfo =
Info
"srv"
"1.0.0"
"go1"
"127.0.0.1"
4222
1024
1
Nothing
Nothing
Nothing
Nothing
Nothing
Nothing
Nothing
Nothing
incrementalInfoParser :: ParserAPI ParsedMessage
incrementalInfoParser = ParserAPI parseInfo
where
parseInfo bytes
| bytes == "INF" =
NeedMore
| bytes == infoFrame =
Emit (ParsedInfo testInfo) ""
| bytes == "PONG\r\n" =
Emit (ParsedPong Pong) ""
| otherwise =
Reject ("unexpected input: " ++ show bytes)
dropPrefixInfoParser :: ParserAPI ParsedMessage
dropPrefixInfoParser = ParserAPI parseInfo
where
parseInfo bytes
| bytes == "XIN" =
DropPrefix 1 "invalid prefix"
| bytes == infoFrame =
Emit (ParsedInfo testInfo) ""
| bytes == "PONG\r\n" =
Emit (ParsedPong Pong) ""
| otherwise =
Reject ("unexpected input: " ++ show bytes)
infoFrame :: BS.ByteString
infoFrame =
"INFO {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n"
newScriptedTransport :: [BS.ByteString] -> TVar [BS.ByteString] -> IO Transport
newScriptedTransport chunks writes = do
remainingChunks <- newTVarIO chunks
pure $
Transport
{ transportRead = \_ ->
atomically $ do
remaining <- readTVar remainingChunks
case remaining of
nextChunk:rest -> do
writeTVar remainingChunks rest
pure nextChunk
[] ->
pure BS.empty
, transportWrite = \bytes ->
atomically $ modifyTVar' writes (<> [bytes])
, transportWriteLazy = \bytes ->
atomically $ modifyTVar' writes (<> [LBS.toStrict bytes])
, transportFlush = pure ()
, transportClose = pure ()
, transportUpgrade = Nothing
}