diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2016 Lars Petersen + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,50 @@ +A pure Haskell MQTT client and server implementation +==================================================== + +## Project goal + +This project aims to supply a rock-solid MQTT 3.1.1 implementation suitable for production use. + +The planned features are: + + - A broker implementation capable of handing and serving several thousands of + connections. + - A client implementation with integrated broker which allows one client to be + used by several threads/consumers simultaneously. + - TLS and WebSocket connections. + - An interface for pluggable authentication and authorization. + - High test and benchmark coverage. + +## Project state (2017-03-22) + + - The broker implementation is nearly feature complete and well-tested. + The [hummingbird](https://github.com/lpeterse/haskell-hummingbird) project + is a full-fledged broker built on-top of this library. + - The client implementation went out of focus for now and is currently + commented out. It's still a planned feature and is essentially a low + hanging fruit as all the protocol parsers etc. are already in place. + +## License + +Permission is hereby granted under the terms of the MIT license: + +> Copyright (c) 2016 Lars Petersen +> +> Permission is hereby granted, free of charge, to any person obtaining +> a copy of this software and associated documentation files (the +> "Software"), to deal in the Software without restriction, including +> without limitation the rights to use, copy, modify, merge, publish, +> distribute, sublicense, and/or sell copies of the Software, and to +> permit persons to whom the Software is furnished to do so, subject to +> the following conditions: +> +> The above copyright notice and this permission notice shall be included +> in all copies or substantial portions of the Software. +> +> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +> EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +> MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +> IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +> CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +> TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +> SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/benchmark/Binary.hs b/benchmark/Binary.hs new file mode 100644 --- /dev/null +++ b/benchmark/Binary.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main ( main, benchmark ) where + +import Criterion.Main +import qualified Data.Binary.Get as BG +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BS +import qualified Data.ByteString.Lazy as BSL + +import Network.MQTT.Topic +import Network.MQTT.Message + +main :: IO () +main = defaultMain benchmark + +benchmark :: [Benchmark] +benchmark = + [ bgroup "Binary" [ + bgroup "ClientConnect (without will)" $ pb $ + ClientConnect "client-identifier" True 65298 Nothing credentials + , bgroup "ClientConnect (with will)" $ pb $ + ClientConnect "client-identifier" True 65298 (Just message) credentials + , bgroup "ClientPublish" $ pb $ + ClientPublish 65298 False message + , bgroup "ClientSubscribe" $ pb $ + ClientSubscribe 65298 $ replicate 23 (filtr, Qos2) + , bgroup "ClientUnsubscribe" $ pb $ + ClientUnsubscribe 65298 $ replicate 23 filtr + ] + ] + where + credentials = Just ("username", Just (Password "password")) + topic = "bla/foo/bar/fnord/gnurp" + filtr = "bla/foo/+/fnord/gnurp/#" + payload = "ashdjahjdahskj hdas dhakjshdas dhakjshdalsjh dashdashd lashdahsad" + message = Message { + msgTopic = topic + , msgBody = payload + , msgQos = Qos1 + , msgRetain = True + } + pb x = [ + -- evaluate the lazy bytestring to strict normal form + bench "build" (nf (BS.toLazyByteString . clientMessageBuilder) x) + , env (pure $ BS.toLazyByteString $ clientMessageBuilder x) $ \bs-> + -- evaluate to weak head normal form (the parser must have touched + -- everything if it could decide whether the input was valid or not) + bench "parse" (whnf (BG.runGet clientMessageParser) bs) + ] diff --git a/benchmark/TopicMatching.hs b/benchmark/TopicMatching.hs new file mode 100644 --- /dev/null +++ b/benchmark/TopicMatching.hs @@ -0,0 +1,28 @@ +module Main where + +import Criterion.Main + +import Data.Bits +import Data.String +import Network.MQTT.Topic +import Network.MQTT.RoutingTree as R + +main :: IO () +main = filterTree `seq` defaultMain [benchmark] + +benchmark :: Benchmark +benchmark = bench "Matching 512 topics against 20000 permissions." $ whnf (foldr (\topic accum-> accum `xor` R.matchTopic topic filterTree) False) topics + +filterTree :: RoutingTree () +filterTree = foldr (\f t-> R.insert (fromString f) () t) mempty filters + +filters :: [String] +filters = [ [x1,x2,x3,'/',x4,x5,x6,'/',x7,x8,x9] | x1<-r,x2<-r,x3<-r,x4<-r,x5<-r,x6<-r,x7<-r,x8<-r,x9<-r] + where + r = ['a'..'c'] + +topics :: [Topic] +topics = fromString <$> f filters + where + f [] = [] + f xs = head xs : f (drop 36 xs) diff --git a/mqtt.cabal b/mqtt.cabal new file mode 100644 --- /dev/null +++ b/mqtt.cabal @@ -0,0 +1,180 @@ +name: mqtt +synopsis: An MQTT protocol implementation. +description: + This project aims to supply a rock-solid MQTT implementation suitable for + production use. +version: 0.1.0.0 +license: MIT +license-file: LICENSE +author: Lars Petersen +maintainer: info@lars-petersen.net +category: Network, IoT +homepage: https://github.com/lpeterse/haskell-mqtt +bug-reports: https://github.com/lpeterse/haskell-mqtt/issues +extra-source-files: README.md +build-type: Simple +cabal-version: >=1.10 + +source-repository head + type: git + location: git://github.com/lpeterse/haskell-mqtt.git + +library + ghc-options: -Wall + exposed-modules: Network.Stack.Server, + Network.MQTT.Client, + Network.MQTT.Message, + Network.MQTT.Trie, + Network.MQTT.Broker, + Network.MQTT.Broker.Authentication, + Network.MQTT.Broker.RetainedMessages, + Network.MQTT.Broker.Server, + Network.MQTT.Broker.Session, + Network.MQTT.Broker.SessionStatistics + other-modules: Control.Concurrent.PrioritySemaphore, + Network.MQTT.Message.QoS, + Network.MQTT.Message.Topic, + Network.MQTT.Broker.Internal + build-depends: async, + attoparsec, + base >=4.8 && <5, + bytestring, + clock, + binary, + exceptions, + text, + containers, + socket, + tls, + uuid, + case-insensitive, + x509, + x509-validation, + websockets, + hslogger + hs-source-dirs: src + default-language: Haskell2010 + +benchmark topic-matching + type: exitcode-stdio-1.0 + ghc-options: -Wall -O2 + default-language: Haskell2010 + hs-source-dirs: benchmark + main-is: TopicMatching.hs + build-depends: + base, + bytestring, + binary, + criterion, + mqtt, + text + +benchmark binary + type: + exitcode-stdio-1.0 + default-language: + Haskell2010 + ghc-options: + -Wall -threaded + hs-source-dirs: + benchmark + main-is: + Binary.hs + build-depends: + base, + bytestring, + binary, + criterion, + mqtt, + text + +test-suite test + type: + exitcode-stdio-1.0 + default-language: + Haskell2010 + ghc-options: + -Wall -fno-warn-orphans + hs-source-dirs: + test + main-is: + Main.hs + other-modules: + BrokerTest, + EncodingTest, + TrieTest, + TopicTest + build-depends: + async + , base >= 4.8 && < 5 + , binary + , tasty >= 0.11 + , tasty-hunit + , tasty-quickcheck + , bytestring + , attoparsec + , exceptions + , mqtt + , text + , tls + , network-uri + , deepseq + , containers + , uuid + +test-suite priority-semaphore + type: + exitcode-stdio-1.0 + default-language: + Haskell2010 + ghc-options: + -Wall + hs-source-dirs: + src, + test + main-is: + PrioritySemaphoreTest.hs + other-modules: + Control.Concurrent.PrioritySemaphore + build-depends: + async + , base >= 4.8 && < 5 + , tasty >= 0.11 + , tasty-hunit + , tasty-quickcheck + , mqtt + +-- This test suite runs with restricted maximum heap size. +-- A memory leak would result in heap usage of more than 4MB +-- leading to test failure. +test-suite retained-store-strictness-test + type: + exitcode-stdio-1.0 + default-language: + Haskell2010 + ghc-options: + -Wall "-with-rtsopts=-M4m" + main-is: + test/RetainedStoreStrictnessTest.hs + build-depends: + base >= 4.8 && <5, + mqtt + +-- This test suite runs with restricted maximum heap size. +-- The routing tree with 1 million nodes and around 3 subscriptions per node +-- took 83MB when measured. +-- We allow for some deviation, but it should not exceed 100MB heap space. +test-suite routing-tree-size-test + type: + exitcode-stdio-1.0 + default-language: + Haskell2010 + ghc-options: + -Wall "-with-rtsopts=-M100m -s" + main-is: + test/TrieSizeTest.hs + build-depends: + base >= 4.8 && <5, + containers, + mqtt, + random diff --git a/src/Control/Concurrent/PrioritySemaphore.hs b/src/Control/Concurrent/PrioritySemaphore.hs new file mode 100644 --- /dev/null +++ b/src/Control/Concurrent/PrioritySemaphore.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE LambdaCase #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Control.Concurrent.PrioritySemaphore +-- Copyright : (c) Lars Petersen 2017 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Control.Concurrent.PrioritySemaphore where + +import Control.Monad ( void ) +import Control.Exception +import Control.Concurrent + +-- | A `PrioritySemaphore` protects a critical section in a way that +-- another thread can interrupt an ongoing computation and take over. +newtype PrioritySemaphore = PrioritySemaphore (MVar (), MVar ThreadId) + +newPrioritySemaphore :: IO PrioritySemaphore +newPrioritySemaphore = PrioritySemaphore <$> ((,) <$> newMVar () <*> newEmptyMVar) + +-- | Enter a critial section with exlusivity and kill any currently running thread. +-- +-- - When no other thread is currently in the section, the action is +-- executed immediately. +-- - When a thread wants to enter and another one is already in the section, +-- the one in the section gets a `ThreadKilled` exception. The waiting +-- thread will enter no earlier than before the killed thread has terminated. +-- This assures that the killed thread can execute cleanup handlers while +-- still having exclusivity wrt to the section. +exclusively :: PrioritySemaphore -> IO a -> IO a +exclusively (PrioritySemaphore (mw, mt)) action = + bracket acquireSemaphore releaseSemaphore (const action) + where + -- Acquiring the semaphore is a two step process in order to assure that + -- only one thread at a time kills the acting thread and waits for its + -- termination. All other threads will wait in front of `mw`. + acquireSemaphore = withMVar mw $ const $ do + -- The thread holds a lock on `mw` here. Find out whether there is + -- thread within the section and eventually kill it. + tryReadMVar mt >>= \case + Nothing -> pure () + Just q -> killThread q + -- The next line will block and unblock as soon as the killed thread + -- has terminated. In case no thread is running, we can put immediately. + -- A race between this line and the line before is ruled out by the + -- the lock on `mw`. + -- As soon as we succeeded in putting in our thread id, we release the + -- lock on `mw` which means we too might get killed while executing + -- our action. + myThreadId >>= putMVar mt + -- Releasing the semaphore means taking the own thread id back + -- from the locking MVar. It is logically guaranteed that the MVar + -- contains the own thread id as a thread will get here only after it has + -- put his thread id in and no other thread can do this as long as the MVar + -- is filled. + releaseSemaphore = const $ void (takeMVar mt) diff --git a/src/Network/MQTT/Broker.hs b/src/Network/MQTT/Broker.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker + ( Broker (brokerAuthenticator) + , newBroker + , publishUpstream + , publishDownstream + , withSession + , getUptime + , getSessions + , getSubscriptions + , lookupSession + ) where + +import Control.Concurrent.MVar +import Control.Concurrent.PrioritySemaphore +import Control.Exception +import Control.Monad +import Data.Int +import qualified Data.IntMap.Strict as IM +import qualified Data.IntSet as IS +import qualified Data.Map.Strict as M +import Data.Maybe +import System.Clock +import qualified System.Log.Logger as Log + +import Network.MQTT.Broker.Authentication +import Network.MQTT.Broker.Internal +import qualified Network.MQTT.Broker.RetainedMessages as RM +import qualified Network.MQTT.Broker.Session as Session +import qualified Network.MQTT.Broker.SessionStatistics as SS +import Network.MQTT.Message +import qualified Network.MQTT.Trie as R + +newBroker :: auth -> IO (Broker auth) +newBroker authenticator = do + now <- sec <$> getTime Realtime + rm <- RM.new + st <-newMVar BrokerState + { brokerMaxSessionIdentifier = SessionIdentifier 0 + , brokerSubscriptions = mempty + , brokerSessions = mempty + , brokerSessionsByPrincipals = mempty + } + pure Broker { + brokerCreatedAt = now + , brokerAuthenticator = authenticator + , brokerRetainedStore = rm + , brokerState = st + } + +withSession :: forall auth. (Authenticator auth) => Broker auth -> ConnectionRequest -> (RejectReason -> IO ()) -> (Session auth -> SessionPresent -> IO ()) -> IO () +withSession broker request sessionRejectHandler sessionAcceptHandler = + (try $ authenticate (brokerAuthenticator broker) request :: IO (Either (AuthenticationException auth) (Maybe PrincipalIdentifier))) >>= \case + Left _ -> sessionRejectHandler ServerUnavailable + Right mp -> case mp of + Nothing -> sessionRejectHandler NotAuthorized + Just principalIdentifier -> bracket + -- In case the principals identity could be determined, we'll either + -- find an associated existing session or create a new one. + -- Getting/creating a session eventually modifies the broker state. + ( getSession broker (principalIdentifier, requestClientIdentifier request) ) + -- This is executed when the current thread terminates (on connection loss). + -- Cleanup actions are executed here (like removing the session when the clean session flag was set). + (\case + Nothing -> pure () + Just (session, _) -> if requestCleanSession request + then Session.terminate session + else Session.reset session + ) + -- This is where the actual connection handler code is invoked. + -- We're using a `PrioritySemaphore` here. This allows other threads for + -- this session to terminate the current one. This is usually the case + -- when the client loses the connection and reconnects, but we have not + -- yet noted the dead connection. The currently running handler thread + -- will receive a `ThreadKilled` exception. + (\case + Nothing -> sessionRejectHandler NotAuthorized + Just (session, sessionPresent)-> + exclusively (sessionSemaphore session) $ do + now <- sec <$> getTime Realtime + let connection = Connection { + connectionCreatedAt = now + , connectionCleanSession = requestCleanSession request + , connectionSecure = requestSecure request + , connectionWebSocket = isJust (requestHttp request) + , connectionRemoteAddress = requestRemoteAddress request + } + bracket_ + ( putMVar (sessionConnection session) connection ) + ( void $ takeMVar (sessionConnection session) ) + ( sessionAcceptHandler session sessionPresent ) + ) + +lookupSession :: SessionIdentifier -> Broker auth -> IO (Maybe (Session auth)) +lookupSession (SessionIdentifier sid) broker = + withMVar (brokerState broker) $ \st-> + pure $ IM.lookup sid (brokerSessions st) + +-- | Either lookup or create a session if none is present (yet). +-- +-- Principal is only looked up initially. Reconnects won't update the +-- permissions etc. Returns Nothing in case the principal identifier cannot +-- be mapped to a principal object. +getSession :: Authenticator auth => Broker auth -> (PrincipalIdentifier, ClientIdentifier) -> IO (Maybe (Session auth, SessionPresent)) +getSession broker pcid@(pid, cid) = + modifyMVar (brokerState broker) $ \st-> + case M.lookup pcid (brokerSessionsByPrincipals st) of + Just (SessionIdentifier sid) -> + case IM.lookup sid (brokerSessions st) of + -- Resuming an existing session.. + Just session -> + pure (st, Just (session, SessionPresent True)) + -- Orphaned session id. This is illegal state. + Nothing -> do + Log.warningM "Broker.getSession" $ "Illegal state: Found orphanded session id " ++ show sid ++ "." + createSession st + -- No session entry found for principal. Creating one. + Nothing -> createSession st + where + createSession st = getPrincipal (brokerAuthenticator broker) pid >>= \case + Nothing -> pure (st, Nothing) + Just principal -> do + now <- sec <$> getTime Realtime + semaphore <- newPrioritySemaphore + subscriptions <- newMVar R.empty + queue <- newMVar (emptyServerQueue $ fromIntegral $ quotaMaxPacketIdentifiers $ principalQuota principal) + queuePending <- newEmptyMVar + mconnection <- newEmptyMVar + mprincipal <- newMVar principal + stats <- SS.new + let SessionIdentifier maxSessionIdentifier = brokerMaxSessionIdentifier st + newSessionIdentifier = maxSessionIdentifier + 1 + newSession = Session + { sessionBroker = broker + , sessionIdentifier = SessionIdentifier newSessionIdentifier + , sessionClientIdentifier = cid + , sessionPrincipalIdentifier = pid + , sessionCreatedAt = now + , sessionConnection = mconnection + , sessionPrincipal = mprincipal + , sessionSemaphore = semaphore + , sessionSubscriptions = subscriptions + , sessionQueue = queue + , sessionQueuePending = queuePending + , sessionStatistics = stats + } + newBrokerState = st + { brokerMaxSessionIdentifier = SessionIdentifier newSessionIdentifier + , brokerSessions = IM.insert newSessionIdentifier newSession (brokerSessions st) + , brokerSessionsByPrincipals = M.insert pcid (SessionIdentifier newSessionIdentifier) (brokerSessionsByPrincipals st) + } + Log.infoM "Broker.createSession" $ "Creating new session with id " ++ show newSessionIdentifier ++ " for " ++ show pid ++ "." + pure (newBrokerState, Just (newSession, SessionPresent False)) + +getUptime :: Broker auth -> IO Int64 +getUptime broker = do + now <- sec <$> getTime Realtime + pure $ now - brokerCreatedAt broker + +getSessions :: Broker auth -> IO (IM.IntMap (Session auth)) +getSessions broker = brokerSessions <$> readMVar (brokerState broker) + +getSubscriptions :: Broker auth -> IO (R.Trie IS.IntSet) +getSubscriptions broker = brokerSubscriptions <$> readMVar (brokerState broker) diff --git a/src/Network/MQTT/Broker/Authentication.hs b/src/Network/MQTT/Broker/Authentication.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/Authentication.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.Authentication +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.Authentication where + +import Control.Exception +import qualified Data.Binary as B +import qualified Data.ByteString as BS +import Data.CaseInsensitive +import Data.UUID as UUID +import Data.Word +import qualified Data.X509 as X509 +import GHC.Generics + +import Network.MQTT.Message +import Network.MQTT.Trie as R + +-- | A peer identity optionally associated with connection/session +-- specific information. +--newtype Principal = Principal T.Text deriving (Eq, Ord, Show) + +-- | An `Authenticator` is able to determine a `Principal`'s identity from a +-- `Request`. +class (Exception (AuthenticationException a)) => Authenticator a where + data AuthenticatorConfig a + -- | This `Exception` may be thrown by any operation within this class. + -- Operations /must/ only throw this type of exception. Other exceptions + -- won't be catched and may kill the broker. + data AuthenticationException a + -- | Create a new authenticator instance from configuration. + newAuthenticator :: AuthenticatorConfig a -> IO a + -- | Try to determine a `Principal`'s identity from a connection `Request`. + -- + -- The operation shall return `Nothing` in case the authentication + -- mechanism is working, but couldn't associate an identity. It shall + -- throw an `AuthenticationException` in case of other problems. + authenticate :: a -> ConnectionRequest -> IO (Maybe PrincipalIdentifier) + -- | Gets a principal by principal primary key (UUID). + -- + -- The operation shall return `Nothing` in case the principal is not / no + -- longer available. It shall throw an `AuthenticationException` in case + -- of other problems. + getPrincipal :: a -> PrincipalIdentifier -> IO (Maybe Principal) + +type PrincipalIdentifier = UUID + +data Principal + = Principal + { principalUsername :: Maybe Username + , principalQuota :: Quota + , principalPublishPermissions :: R.Trie () + , principalSubscribePermissions :: R.Trie () + , principalRetainPermissions :: R.Trie () + } deriving (Eq, Show, Generic) + +data Quota + = Quota + { quotaMaxIdleSessionTTL :: Word64 + , quotaMaxPacketSize :: Word64 + , quotaMaxPacketIdentifiers :: Word64 + , quotaMaxQueueSizeQoS0 :: Word64 + , quotaMaxQueueSizeQoS1 :: Word64 + , quotaMaxQueueSizeQoS2 :: Word64 + } deriving (Eq, Ord, Show, Generic) + +instance B.Binary Quota +instance B.Binary Principal + +-- | This class defines how the information gathered from a +-- connection request looks like. An `Authenticator` may use +-- whatever information it finds suitable to authenticate the `Principal`. +data ConnectionRequest + = ConnectionRequest + { requestClientIdentifier :: ClientIdentifier, + requestCleanSession :: Bool, + -- | Is this connection secure in terms of + -- [Transport Layer Security](https://en.wikipedia.org/wiki/Transport_Layer_Security)? + requestSecure :: Bool, + -- | The username and password supplied with the MQTT handshake. + requestCredentials :: Maybe (Username, Maybe Password), + -- | The HTTP request head in case the client connected via + -- [WebSocket](https://en.wikipedia.org/wiki/WebSocket). + requestHttp :: Maybe (BS.ByteString, [(CI BS.ByteString, BS.ByteString)]), + -- | An [X.509 certificate](https://en.wikipedia.org/wiki/X.509) chain + -- supplied by the peer. + -- It can be assumed that the transport layer implementation already + -- verified that the peer owns the corresponding private key. The validation + -- of the certificate claims (including certificate chain checking) /must/ + -- be performed by the `Authenticator`. + requestCertificateChain :: Maybe X509.CertificateChain, + requestRemoteAddress :: Maybe BS.ByteString + } deriving (Show) diff --git a/src/Network/MQTT/Broker/Internal.hs b/src/Network/MQTT/Broker/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/Internal.hs @@ -0,0 +1,219 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.Internal +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.Internal where + +import Control.Concurrent.MVar +import Control.Concurrent.PrioritySemaphore +import Control.Monad +import qualified Data.Binary as B +import qualified Data.ByteString as BS +import Data.Int +import qualified Data.IntMap.Strict as IM +import qualified Data.IntSet as IS +import qualified Data.Map.Strict as M +import qualified Data.Sequence as Seq +import GHC.Generics (Generic) + +import Network.MQTT.Broker.Authentication hiding (getPrincipal) +import qualified Network.MQTT.Broker.RetainedMessages as RM +import qualified Network.MQTT.Broker.SessionStatistics as SS +import Network.MQTT.Message +import qualified Network.MQTT.Trie as R + +data Broker auth + = Broker + { brokerCreatedAt :: Int64 + , brokerAuthenticator :: auth + , brokerRetainedStore :: RM.RetainedStore + , brokerState :: MVar (BrokerState auth) + } + +data BrokerState auth + = BrokerState + { brokerMaxSessionIdentifier :: !SessionIdentifier + , brokerSubscriptions :: !(R.Trie IS.IntSet) + , brokerSessions :: !(IM.IntMap (Session auth)) + , brokerSessionsByPrincipals :: !(M.Map (PrincipalIdentifier, ClientIdentifier) SessionIdentifier) + } + +newtype SessionIdentifier = SessionIdentifier Int deriving (Eq, Ord, Show, Enum, Generic) + +data Session auth + = Session + { sessionBroker :: !(Broker auth) + , sessionIdentifier :: !SessionIdentifier + , sessionClientIdentifier :: !ClientIdentifier + , sessionPrincipalIdentifier :: !PrincipalIdentifier + , sessionCreatedAt :: !Int64 + , sessionConnection :: !(MVar Connection) + , sessionPrincipal :: !(MVar Principal) + , sessionSemaphore :: !PrioritySemaphore + , sessionSubscriptions :: !(MVar (R.Trie QoS)) + , sessionQueue :: !(MVar ServerQueue) + , sessionQueuePending :: !(MVar ()) + , sessionStatistics :: !SS.Statistics + } + +data Connection + = Connection + { connectionCreatedAt :: !Int64 + , connectionCleanSession :: !Bool + , connectionSecure :: !Bool + , connectionWebSocket :: !Bool + , connectionRemoteAddress :: !(Maybe BS.ByteString) + } deriving (Eq, Ord, Show, Generic) + +data ServerQueue + = ServerQueue + { queuePids :: !(Seq.Seq PacketIdentifier) + , outputBuffer :: !(Seq.Seq ServerPacket) + , queueQoS0 :: !(Seq.Seq Message) + , queueQoS1 :: !(Seq.Seq Message) + , queueQoS2 :: !(Seq.Seq Message) + , notAcknowledged :: !(IM.IntMap Message) -- We sent a `QoS1` message and have not yet received the @PUBACK@. + , notReceived :: !(IM.IntMap Message) -- We sent a `QoS2` message and have not yet received the @PUBREC@. + , notReleased :: !(IM.IntMap Message) -- We received as `QoS2` message, sent the @PUBREC@ and wait for the @PUBREL@. + , notComplete :: !IS.IntSet -- We sent a @PUBREL@ and have not yet received the @PUBCOMP@. + } + +instance B.Binary SessionIdentifier +instance B.Binary Connection + +instance Eq (Session auth) where + (==) s1 s2 = (==) (sessionIdentifier s1) (sessionIdentifier s2) + +instance Ord (Session auth) where + compare s1 s2 = compare (sessionIdentifier s1) (sessionIdentifier s2) + +instance Show (Session auth) where + show session = + "Session { identifier = " ++ show (sessionIdentifier session) + ++ ", principal = " ++ show (sessionPrincipalIdentifier session) + ++ ", client = " ++ show (sessionClientIdentifier session) ++ " }" + +-- | Inject a message downstream into the broker. It will be delivered +-- to all subscribed sessions within this broker instance. +publishDownstream :: Broker auth -> Message -> IO () +publishDownstream broker msg = do + RM.store msg (brokerRetainedStore broker) + let topic = msgTopic msg + st <- readMVar (brokerState broker) + forM_ (IS.elems $ R.lookup topic $ brokerSubscriptions st) $ \key-> + case IM.lookup (key :: Int) (brokerSessions st) of + Nothing -> + putStrLn "WARNING: dead session reference" + Just session -> publishMessage session msg + +-- | Publish a message upstream on the broker. +-- +-- * As long as the broker is not clustered upstream=downstream. +-- * FUTURE NOTE: In clustering mode this shall distribute the message +-- to other brokers or upwards when the brokers form a hierarchy. +publishUpstream :: Broker auth -> Message -> IO () +publishUpstream = publishDownstream + +notePending :: Session auth -> IO () +notePending = void . flip tryPutMVar () . sessionQueuePending + +waitPending :: Session auth -> IO () +waitPending = void . readMVar . sessionQueuePending + +emptyServerQueue :: Int -> ServerQueue +emptyServerQueue i = ServerQueue + { queuePids = Seq.fromList $ fmap PacketIdentifier [0 .. min i 65535] + , outputBuffer = mempty + , queueQoS0 = mempty + , queueQoS1 = mempty + , queueQoS2 = mempty + , notAcknowledged = mempty + , notReceived = mempty + , notReleased = mempty + , notComplete = mempty + } + +publishMessage :: Session auth -> Message -> IO () +publishMessage session msg = do + subscriptions <- readMVar (sessionSubscriptions session) + case R.findMaxBounded (msgTopic msg) subscriptions of + Nothing -> pure () + Just qos -> enqueueMessage session msg { msgQoS = qos } + +publishMessages :: Foldable t => Session auth -> t Message -> IO () +publishMessages session msgs = + forM_ msgs (publishMessage session) + +-- | This enqueues a message for transmission to the client. +-- +-- * This operations eventually terminates the session on queue overflow. +-- The caller will not notice this and the operation will not throw an exception. +enqueueMessage :: Session auth -> Message -> IO () +enqueueMessage session msg = do + quota <- principalQuota <$> readMVar (sessionPrincipal session) + success <- modifyMVar (sessionQueue session) $ \queue-> case msgQoS msg of + QoS0 -> if (fromIntegral $ quotaMaxQueueSizeQoS0 quota) > Seq.length (queueQoS0 queue) + then pure $ (, True) $! queue { queueQoS0 = queueQoS0 queue Seq.|> msg } + else pure $ (, True) $! queue { queueQoS0 = Seq.drop 1 $ queueQoS0 queue Seq.|> msg } + QoS1 -> if (fromIntegral $ quotaMaxQueueSizeQoS1 quota) > Seq.length (queueQoS1 queue) + then pure $ (, True) $! queue { queueQoS1 = queueQoS1 queue Seq.|> msg } + else pure (queue, False) + QoS2 -> if (fromIntegral $ quotaMaxQueueSizeQoS2 quota) > Seq.length (queueQoS2 queue) + then pure $ (, True) $! queue { queueQoS2 = queueQoS2 queue Seq.|> msg } + else pure (queue, False) + if success + -- Notify the sending thread that something has been enqueued! + then notePending session + -- Kill the session. + else terminate session + +-- TODO: make more efficient +enqueueMessages :: Foldable t => Session auth -> t Message -> IO () +enqueueMessages session msgs = + forM_ msgs (enqueueMessage session) + +-- | Terminate a session. +-- +-- * An eventually connected client gets disconnected. +-- * The session subscriptions are removed from the subscription tree +-- which means that it will receive no more messages. +-- * The session will be unlinked from the broker which means +-- that clients cannot resume it anymore under this client identifier. +terminate :: Session auth -> IO () +terminate session = + -- This assures that the client gets disconnected. The code is executed + -- _after_ the current client handler that has terminated. + -- TODO Race: New clients may try to connect while we are in here. + -- This would not make the state inconsistent, but kill this thread. + -- What we need is another `exclusivelyUninterruptible` function for + -- the priority semaphore. + exclusively (sessionSemaphore session) $ + modifyMVarMasked_ (brokerState $ sessionBroker session) $ \st-> + withMVarMasked (sessionSubscriptions session) $ \subscriptions-> + pure st + { brokerSessions = IM.delete sid ( brokerSessions st ) + -- Remove the session id from the (principal, client) -> sessionid + -- mapping. Remove empty leaves in this mapping, too. + , brokerSessionsByPrincipals = M.delete + ( sessionPrincipalIdentifier session + , sessionClientIdentifier session) + (brokerSessionsByPrincipals st) + -- Remove the session id from each set that the session + -- subscription tree has a corresponding value for (which is ignored). + , brokerSubscriptions = R.differenceWith + (\b _-> Just (IS.delete sid b) ) + ( brokerSubscriptions st ) subscriptions + } + where + SessionIdentifier sid = sessionIdentifier session diff --git a/src/Network/MQTT/Broker/RetainedMessages.hs b/src/Network/MQTT/Broker/RetainedMessages.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/RetainedMessages.hs @@ -0,0 +1,110 @@ +{-# LANGUAGE OverloadedStrings #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.RetainedMessages +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.RetainedMessages where + +import Control.Applicative hiding (empty) +import Control.Concurrent.MVar +import qualified Data.ByteString.Lazy as BSL +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Prelude hiding (null) + +import qualified Network.MQTT.Message as Message +import qualified Network.MQTT.Message.Topic as Topic + +newtype RetainedStore = RetainedStore { unstore :: MVar RetainedTree } +newtype RetainedTree = RetainedTree { untree :: M.Map Topic.Level RetainedNode } +data RetainedNode = RetainedNode !RetainedTree !(Maybe Message.Message) + +new :: IO RetainedStore +new = RetainedStore <$> newMVar empty + +store :: Message.Message -> RetainedStore -> IO () +store msg (RetainedStore mvar) + | retain = modifyMVar_ mvar $ \tree-> + -- The seq ($!) is important for not leaking memory! + pure $! if BSL.null body + then delete msg tree + else insert msg tree + | otherwise = pure () + where + Message.Payload body = Message.msgPayload msg + Message.Retain retain = Message.msgRetain msg + +retrieve :: Topic.Filter -> RetainedStore -> IO (S.Set Message.Message) +retrieve filtr (RetainedStore mvar) = + lookupFilter filtr <$> readMVar mvar + +empty :: RetainedTree +empty = RetainedTree mempty + +null :: RetainedTree -> Bool +null = M.null . untree + +insert :: Message.Message -> RetainedTree -> RetainedTree +insert msg = union (singleton msg) + +delete :: Message.Message -> RetainedTree -> RetainedTree +delete msg = difference (singleton msg) + +singleton :: Message.Message -> RetainedTree +singleton msg = + let l :| ls = Topic.topicLevels (Message.msgTopic msg) + in RetainedTree (M.singleton l $ node ls) + where + node [] = RetainedNode empty $! Just $! msg + node (x:xs) = RetainedNode (RetainedTree $ M.singleton x $ node xs) Nothing + +-- | The expression `union t1 t2` takes the left-biased union of `t1` and `t2`. +union :: RetainedTree -> RetainedTree -> RetainedTree +union (RetainedTree m1) (RetainedTree m2) = + RetainedTree $ M.unionWith merge m1 m2 + where + merge (RetainedNode t1 mm1) (RetainedNode t2 mm2) = + RetainedNode (t1 `union` t2) $! case mm1 <|> mm2 of + Nothing -> Nothing + Just mm -> Just $! mm + +difference :: RetainedTree -> RetainedTree -> RetainedTree +difference (RetainedTree m1) (RetainedTree m2) = + RetainedTree $ M.differenceWith diff m1 m2 + where + diff (RetainedNode t1 mm1) (RetainedNode t2 mm2) + | null t3 && isNothing mm3 = Nothing + | otherwise = Just (RetainedNode t3 mm3) + where + t3 = difference t1 t2 + mm3 = case mm2 of + Just _ -> Nothing + Nothing -> mm1 + +lookupFilter :: Topic.Filter -> RetainedTree -> S.Set Message.Message +lookupFilter filtr t = + let l :| ls = Topic.filterLevels filtr in collect l ls t + where + collect l ls tree@(RetainedTree m) = case l of + "#" -> allTree tree + "+" -> M.foldl (\s node-> s `S.union` pathNode ls node) S.empty m + _ -> fromMaybe S.empty $ pathNode ls <$> M.lookup l m + allTree (RetainedTree branches) = + M.foldl (\s node-> s `S.union` allNode node) S.empty branches + allNode (RetainedNode subtree mmsg) = + case mmsg of + Nothing -> allTree subtree + Just msg -> S.insert msg (allTree subtree) + pathNode [] (RetainedNode _ mmsg) = + fromMaybe S.empty $ S.singleton <$> mmsg + pathNode (x:xs) (RetainedNode subtree mmsg) = + case x of + "#"-> fromMaybe id (S.insert <$> mmsg) (collect x xs subtree) + _ -> collect x xs subtree diff --git a/src/Network/MQTT/Broker/Server.hs b/src/Network/MQTT/Broker/Server.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/Server.hs @@ -0,0 +1,296 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.Server +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.Server + ( serveConnection + , MQTT () + , MqttServerTransportStack (..) + , SS.Server ( .. ) + , SS.ServerConfig ( .. ) + , SS.ServerConnection ( .. ) + , SS.ServerException ( .. ) + ) where + +import Control.Concurrent +import Control.Concurrent.Async +import qualified Control.Exception as E +import Control.Monad +import qualified Data.Binary.Get as SG +import qualified Data.ByteString as BS +import Data.Int +import Data.IORef +import Data.Typeable +import qualified Network.Stack.Server as SS +import qualified Network.WebSockets as WS +import qualified System.Log.Logger as Log +import qualified System.Socket as S + +import Network.MQTT.Message +import Network.MQTT.Broker.Authentication +import qualified Network.MQTT.Broker.Internal as Session +import qualified Network.MQTT.Broker as Broker +import qualified Network.MQTT.Broker.Session as Session + +instance (Typeable transport) => E.Exception (SS.ServerException (MQTT transport)) + +data MQTT transport + +class SS.ServerStack a => MqttServerTransportStack a where + getConnectionRequest :: SS.ServerConnectionInfo a -> IO ConnectionRequest + +instance (Typeable f, Typeable t, Typeable p, S.Family f, S.Protocol p, S.Type t, S.HasNameInfo f) => MqttServerTransportStack (S.Socket f t p) where + getConnectionRequest (SS.SocketServerConnectionInfo addr) = do + remoteAddr <- S.hostName <$> S.getNameInfo addr (S.niNumericHost `mappend` S.niNumericService) + pure ConnectionRequest { + requestClientIdentifier = ClientIdentifier mempty + , requestSecure = False + , requestCleanSession = True + , requestCredentials = Nothing + , requestHttp = Nothing + , requestCertificateChain = Nothing + , requestRemoteAddress = Just remoteAddr + } + +instance (SS.StreamServerStack a, MqttServerTransportStack a) => MqttServerTransportStack (SS.WebSocket a) where + getConnectionRequest (SS.WebSocketServerConnectionInfo tci rh) = do + req <- getConnectionRequest tci + pure req { + requestHttp = Just (WS.requestPath rh, WS.requestHeaders rh) + } + +instance (SS.StreamServerStack a, MqttServerTransportStack a) => MqttServerTransportStack (SS.TLS a) where + getConnectionRequest (SS.TlsServerConnectionInfo tci mcc) = do + req <- getConnectionRequest tci + pure req { + requestSecure = True + , requestCertificateChain = mcc + } + +instance (SS.StreamServerStack transport) => SS.ServerStack (MQTT transport) where + data Server (MQTT transport) = MqttServer + { mqttTransportServer :: SS.Server transport + } + data ServerConfig (MQTT transport) = MqttServerConfig + { mqttTransportConfig :: SS.ServerConfig transport + } + data ServerConnection (MQTT transport) = MqttServerConnection + { mqttTransportConnection :: SS.ServerConnection transport + , mqttTransportLeftover :: MVar BS.ByteString + } + data ServerConnectionInfo (MQTT transport) = MqttServerConnectionInfo + { mqttTransportServerConnectionInfo :: SS.ServerConnectionInfo transport + } + data ServerException (MQTT transport) + = ProtocolViolation String + | MessageTooLong + | ConnectionRejected RejectReason + | KeepAliveTimeoutException + deriving (Eq, Ord, Show, Typeable) + withServer config handle = + SS.withServer (mqttTransportConfig config) $ \server-> + handle (MqttServer server) + withConnection server handler = + SS.withConnection (mqttTransportServer server) $ \connection info-> + flip handler (MqttServerConnectionInfo info) =<< MqttServerConnection + <$> pure connection + <*> newMVar mempty + +-- TODO: eventually too strict with message size tracking +instance (SS.StreamServerStack transport) => SS.MessageServerStack (MQTT transport) where + type ClientMessage (MQTT transport) = ClientPacket + type ServerMessage (MQTT transport) = ServerPacket + sendMessage connection = + SS.sendStreamBuilder (mqttTransportConnection connection) 8192 . serverPacketBuilder + sendMessages connection msgs = + SS.sendStreamBuilder (mqttTransportConnection connection) 8192 $ foldl (\b m-> b `mappend` serverPacketBuilder m) mempty msgs + receiveMessage connection maxMsgSize = + modifyMVar (mqttTransportLeftover connection) (execute 0 . SG.pushChunk decode) + where + fetch = SS.receiveStream (mqttTransportConnection connection) 4096 + decode = SG.runGetIncremental clientPacketParser + execute received result + | received > maxMsgSize = E.throwIO (MessageTooLong :: SS.ServerException (MQTT transport)) + | otherwise = case result of + SG.Partial continuation -> do + bs <- fetch + if BS.null bs + then execute received (continuation Nothing) + else execute (received + fromIntegral (BS.length bs)) (continuation $ Just bs) + SG.Fail _ _ failure -> + E.throwIO (ProtocolViolation failure :: SS.ServerException (MQTT transport)) + SG.Done leftover' _ msg -> + pure (leftover', msg) + consumeMessages connection maxMsgSize consume = + modifyMVar_ (mqttTransportLeftover connection) (execute 0 . SG.pushChunk decode) + where + fetch = SS.receiveStream (mqttTransportConnection connection) 4096 + decode = SG.runGetIncremental clientPacketParser + execute received result + | received > maxMsgSize = E.throwIO (MessageTooLong :: SS.ServerException (MQTT transport)) + | otherwise = case result of + SG.Partial continuation -> do + bs <- fetch + if BS.null bs + then execute received (continuation Nothing) + else execute (received + fromIntegral (BS.length bs)) (continuation $ Just bs) + SG.Fail _ _ failure -> + E.throwIO (ProtocolViolation failure :: SS.ServerException (MQTT transport)) + SG.Done leftover' _ msg -> do + done <- consume msg + if done + then pure leftover' + else execute 0 (SG.pushChunk decode leftover') + +deriving instance Show (SS.ServerConnectionInfo transport) => Show (SS.ServerConnectionInfo (MQTT transport)) + +serveConnection :: forall transport auth. (SS.StreamServerStack transport, MqttServerTransportStack transport, Authenticator auth) => Broker.Broker auth -> SS.ServerConnection (MQTT transport) -> SS.ServerConnectionInfo (MQTT transport) -> IO () +serveConnection broker conn connInfo = do + recentActivity <- newIORef True + req <- getConnectionRequest (mqttTransportServerConnectionInfo connInfo) + msg <- SS.receiveMessage conn maxInitialPacketSize + case msg of + ClientConnectUnsupported -> do + Log.warningM "Server.connection" $ "Connection from " + ++ show (requestRemoteAddress req) ++ " rejected: UnacceptableProtocolVersion" + void $ SS.sendMessage conn (ServerConnectionRejected UnacceptableProtocolVersion) + -- Communication ends here gracefully. The caller shall close the connection. + ClientConnect {} -> do + let -- | This one is called when the authenticator decided to reject the request. + sessionRejectHandler reason = do + Log.warningM "Server.connection" $ "Connection rejected: " ++ show reason + void $ SS.sendMessage conn (ServerConnectionRejected reason) + -- Communication ends here gracefully. The caller shall close the connection. + + -- | This part is where the threads for a connection are created + -- (one for input, one for output and one watchdog thread). + sessionAcceptHandler session sessionPresent@(SessionPresent sp) = do + principal <- Session.getPrincipal session + Log.infoM "Server.connection" $ "Connection accepted: Associated " + ++ show principal ++ (if sp then " with existing session " + ++ show (Session.sessionIdentifier session) ++ "." else " with new session.") + void $ SS.sendMessage conn (ServerConnectionAccepted sessionPresent) + foldl1 race_ + [ handleInput recentActivity session + , handleOutput session + , keepAlive recentActivity (connectKeepAlive msg) session + ] `E.catch` (\e-> do + Log.warningM "Server.connection" $"Session " ++ show (Session.sessionIdentifier session) + ++ ": Connection terminated with exception: " ++ show (e :: E.SomeException) + E.throwIO e + ) + Log.infoM "Server.connection" $ + "Session " ++ show (Session.sessionIdentifier session) ++ ": Graceful disconnect." + -- Extend the request object with information gathered from the connect packet. + request = req { + requestClientIdentifier = connectClientIdentifier msg + , requestCleanSession = cleanSession + , requestCredentials = connectCredentials msg + } + where + CleanSession cleanSession = connectCleanSession msg + Log.infoM "Server.connection" $ "Connection request: " ++ show request + Broker.withSession broker request sessionRejectHandler sessionAcceptHandler + _ -> pure () -- TODO: Don't parse not-CONN packets in the first place! + where + -- The size of the initial CONN packet shall somewhat be limited to a moderate size. + -- This value is assumed to make no problems while still protecting + -- the servers resources against exaustion attacks. + maxInitialPacketSize :: Int64 + maxInitialPacketSize = 65535 + -- The keep alive thread wakes up every `keepAlive/2` seconds. + -- When it detects no recent activity, it sleeps one more full `keepAlive` + -- interval and checks again. When it still finds no recent activity, it + -- throws an exception. + -- That way a timeout will be detected between 1.5 and 2 `keep alive` + -- intervals after the last actual client activity. + keepAlive :: IORef Bool -> KeepAliveInterval -> Session.Session auth -> IO () + keepAlive recentActivity (KeepAliveInterval interval) session = forever $ do + writeIORef recentActivity False + threadDelay regularInterval + activity <- readIORef recentActivity + unless activity $ do + threadDelay regularInterval + activity' <- readIORef recentActivity + unless activity' $ do + -- Alert state: The client must get active within the next interval. + Log.warningM "Server.connection.keepAlive" $ "Session " ++ show (Session.sessionIdentifier session) ++ ": Client is overdue." + threadDelay regularInterval + activity'' <- readIORef recentActivity + unless activity'' $ E.throwIO (KeepAliveTimeoutException :: SS.ServerException (MQTT transport)) + where + regularInterval = fromIntegral interval * 500000 + -- | This thread is responsible for continuously processing input. + -- It blocks on reading the input stream until input becomes available. + -- Input is consumed no faster than it can be processed. + -- + -- * Read packet from the input stream. + -- * Note that there was activity (TODO: a timeout may occur when a packet + -- takes too long to transmit due to its size). + -- * Process and dispatch the message internally. + -- * Repeat and eventually wait again. + -- * Eventually throws `ServerException`s. + handleInput :: IORef Bool -> Session.Session auth -> IO () + handleInput recentActivity session = do + maxPacketSize <- fromIntegral . quotaMaxPacketSize . principalQuota <$> Session.getPrincipal session + SS.consumeMessages conn maxPacketSize $ \packet-> do + writeIORef recentActivity True + --Log.debugM "Server.connection.handleInput" $ take 50 $ show packet + case packet of + ClientConnect {} -> + E.throwIO (ProtocolViolation "Unexpected CONN packet." :: SS.ServerException (MQTT transport)) + ClientConnectUnsupported -> + E.throwIO (ProtocolViolation "Unexpected CONN packet (of unsupported protocol version)." :: SS.ServerException (MQTT transport)) + ClientPublish pid dup msg -> do + Session.processPublish session pid dup msg + pure False + ClientPublishAcknowledged pid -> do + Session.processPublishAcknowledged session pid + pure False + ClientPublishReceived pid -> do + Session.processPublishReceived session pid + pure False + ClientPublishRelease pid -> do + Session.processPublishRelease session pid + pure False + ClientPublishComplete pid -> do + Session.processPublishComplete session pid + pure False + ClientSubscribe pid filters -> do + Session.subscribe session pid filters + pure False + ClientUnsubscribe pid filters -> do + Session.unsubscribe session pid filters + pure False + ClientPingRequest -> do + Session.enqueuePingResponse session + pure False + ClientDisconnect -> + pure True + -- | This thread is responsible for continuously transmitting to the client + -- and reading from the output queue. + -- + -- * It blocks on Session.waitPending until output gets available. + -- * When output is available, it fetches a whole sequence of messages + -- from the output queue. + -- * It then uses the optimized SS.sendMessages operation which fills + -- a whole chunk with as many messages as possible and sends the chunks + -- each with a single system call. This is _very_ important for high + -- throughput. + -- * Afterwards, it repeats and eventually waits again. + handleOutput :: Session.Session auth -> IO () + handleOutput session = forever $ do + -- The `waitPending` operation is blocking until messages get available. + Session.waitPending session + msgs <- Session.dequeue session + SS.sendMessages conn msgs diff --git a/src/Network/MQTT/Broker/Session.hs b/src/Network/MQTT/Broker/Session.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/Session.hs @@ -0,0 +1,334 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.Session +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.Session + ( publish + , subscribe + , unsubscribe + , disconnect + , terminate + + , getSubscriptions + , getConnection + , getPrincipal + , getFreePacketIdentifiers + + -- TODO: private + , reset + , notePending + , waitPending + , publishMessage + , publishMessages + , enqueuePingResponse + , enqueueMessage + , enqueueSubscribeAcknowledged + , enqueueUnsubscribeAcknowledged + , dequeue + , processPublish + , processPublishRelease + , processPublishReceived + , processPublishComplete + , processPublishAcknowledged + , Session (..) + , SessionIdentifier (..) + , Connection (..) + ) where + +import Control.Concurrent.MVar +import Control.Concurrent.PrioritySemaphore +import Control.Monad +import Data.Bool +import Data.Functor.Identity +import qualified Data.IntMap as IM +import qualified Data.IntSet as IS +import Data.Maybe +import Data.Monoid +import qualified Data.Sequence as Seq + +import Network.MQTT.Broker.Authentication hiding (getPrincipal) +import Network.MQTT.Broker.Internal +import qualified Network.MQTT.Broker.RetainedMessages as RM +import qualified Network.MQTT.Broker.SessionStatistics as SS +import Network.MQTT.Message +import qualified Network.MQTT.Trie as R + +publish :: Session auth -> Message -> IO () +publish session msg = do + principal <- readMVar (sessionPrincipal session) + -- A topic is permitted if it yields a match in the publish permission tree. + if R.matchTopic (msgTopic msg) (principalPublishPermissions principal) + then do + if retain && R.matchTopic (msgTopic msg) (principalRetainPermissions principal) + then do + RM.store msg (brokerRetainedStore $ sessionBroker session) + SS.accountRetentionsAccepted stats 1 + else + SS.accountRetentionsDropped stats 1 + publishUpstream (sessionBroker session) msg + SS.accountPublicationsAccepted stats 1 + else + SS.accountPublicationsDropped stats 1 + where + stats = sessionStatistics session + Retain retain = msgRetain msg + +subscribe :: Session auth -> PacketIdentifier -> [(Filter, QoS)] -> IO () +subscribe session pid filters = do + principal <- readMVar (sessionPrincipal session) + checkedFilters <- mapM (checkPermission principal) filters + let subscribeFilters = mapMaybe (\(filtr,mqos)->(filtr,) <$> mqos) checkedFilters + qosTree = R.insertFoldable subscribeFilters R.empty + sidTree = R.map (const $ IS.singleton sid) qosTree + -- Do the accounting for the session statistics. + -- TODO: Do this as a transaction below. + let countAccepted = length subscribeFilters + let countDenied = length filters - countAccepted + SS.accountSubscriptionsAccepted (sessionStatistics session) $ fromIntegral countAccepted + SS.accountSubscriptionsDenied (sessionStatistics session) $ fromIntegral countDenied + -- Force the `qosTree` in order to lock the broker as little as possible. + -- The `sidTree` is left lazy. + qosTree `seq` do + modifyMVarMasked_ (brokerState $ sessionBroker session) $ \bst-> do + modifyMVarMasked_ + ( sessionSubscriptions session ) + ( pure . R.unionWith max qosTree ) + pure $ bst { brokerSubscriptions = R.unionWith IS.union (brokerSubscriptions bst) sidTree } + enqueueSubscribeAcknowledged session pid (fmap snd checkedFilters) + forM_ checkedFilters $ \(filtr,_qos)-> + publishMessages session =<< RM.retrieve filtr (brokerRetainedStore $ sessionBroker session) + where + SessionIdentifier sid = sessionIdentifier session + checkPermission principal (filtr, qos) = do + let isPermitted = R.matchFilter filtr (principalSubscribePermissions principal) + pure (filtr, if isPermitted then Just qos else Nothing) + +unsubscribe :: Session auth -> PacketIdentifier -> [Filter] -> IO () +unsubscribe session pid filters = + -- Force the `unsubBrokerTree` first in order to lock the broker as little as possible. + unsubBrokerTree `seq` do + modifyMVarMasked_ (brokerState $ sessionBroker session) $ \bst-> do + modifyMVarMasked_ + ( sessionSubscriptions session ) + ( pure . flip (R.differenceWith (const . const Nothing)) unsubBrokerTree ) + pure $ bst { brokerSubscriptions = R.differenceWith + (\is (Identity i)-> Just (IS.delete i is)) + (brokerSubscriptions bst) unsubBrokerTree } + enqueueUnsubscribeAcknowledged session pid + where + SessionIdentifier sid = sessionIdentifier session + unsubBrokerTree = R.insertFoldable ( fmap (,Identity sid) filters ) R.empty + +-- | Disconnect a session. +disconnect :: Session auth -> IO () +disconnect session = + -- This assures that the client gets disconnected by interrupting + -- the current client handler thread (if any). + exclusively (sessionSemaphore session) (pure ()) + +-- | Reset the session state after a reconnect. +-- +-- * All output buffers will be cleared. +-- * Output buffers will be filled with retransmissions. +reset :: Session auth -> IO () +reset session = + modifyMVar_ (sessionQueue session) (\q-> pure $! resetQueue q) + +-- | Enqueue a PINGRESP to be sent as soon as the output thread is available. +-- +-- The PINGRESP will be inserted with highest priority in front of all other enqueued messages. +enqueuePingResponse :: Session auth -> IO () +enqueuePingResponse session = do + modifyMVar_ (sessionQueue session) $ \queue-> + pure $! queue { outputBuffer = ServerPingResponse Seq.<| outputBuffer queue } + -- IMPORTANT: Notify the sending thread that something has been enqueued! + notePending session + +enqueueSubscribeAcknowledged :: Session auth -> PacketIdentifier -> [Maybe QoS] -> IO () +enqueueSubscribeAcknowledged session pid mqoss = do + modifyMVar_ (sessionQueue session) $ \queue-> + pure $! queue { outputBuffer = outputBuffer queue Seq.|> ServerSubscribeAcknowledged pid mqoss } + notePending session + +enqueueUnsubscribeAcknowledged :: Session auth -> PacketIdentifier -> IO () +enqueueUnsubscribeAcknowledged session pid = do + modifyMVar_ (sessionQueue session) $ \queue-> + pure $! queue { outputBuffer = outputBuffer queue Seq.|> ServerUnsubscribeAcknowledged pid} + notePending session + +-- | Blocks until messages are available and prefers non-qos0 messages over +-- qos0 messages. +dequeue :: Session auth -> IO (Seq.Seq ServerPacket) +dequeue session = + modifyMVar (sessionQueue session) $ \queue-> do + let q = normalizeQueue queue + if | not (Seq.null $ outputBuffer q) -> pure (q { outputBuffer = mempty }, outputBuffer q) + | not (Seq.null $ queueQoS0 q) -> pure (q { queueQoS0 = mempty }, fmap (ServerPublish (PacketIdentifier (-1)) (Duplicate False)) (queueQoS0 q)) + | otherwise -> clearPending >> pure (q, mempty) + where + -- | In case all queues are empty, we need to clear the `pending` variable. + -- ATTENTION: An implementation error would cause the `dequeue` operation + -- to rush right through the blocking call leading to enourmous CPU usage. + clearPending :: IO () + clearPending = void $ tryTakeMVar (sessionQueuePending session) + +-- | Process a @PUB@ message received from the peer. +-- +-- Different handling depending on message qos. +processPublish :: Session auth -> PacketIdentifier -> Duplicate -> Message -> IO () +processPublish session pid@(PacketIdentifier p) _dup msg = + case msgQoS msg of + QoS0 -> + publish session msg + QoS1 -> do + publish session msg + modifyMVar_ (sessionQueue session) $ \q-> pure $! q { + outputBuffer = outputBuffer q Seq.|> ServerPublishAcknowledged pid + } + notePending session + QoS2 -> do + modifyMVar_ (sessionQueue session) $ \q-> pure $! q { + outputBuffer = outputBuffer q Seq.|> ServerPublishReceived pid + , notReleased = IM.insert p msg (notReleased q) + } + notePending session + +-- | Note that a QoS1 message has been received by the peer. +-- +-- This shall be called when a @PUBACK@ has been received by the peer. +-- We release the message from our buffer and the transaction is then complete. +processPublishAcknowledged :: Session auth -> PacketIdentifier -> IO () +processPublishAcknowledged session (PacketIdentifier pid) = do + modifyMVar_ (sessionQueue session) $ \q-> pure $! q { + -- The packet identifier is free for reuse only if it actually was in the set of notAcknowledged messages. + queuePids = bool (queuePids q) (PacketIdentifier pid Seq.<| queuePids q) (IM.member pid (notAcknowledged q)) + , notAcknowledged = IM.delete pid (notAcknowledged q) + } + -- See code of `processPublishComplete` for explanation. + notePending session + +-- | Note that a QoS2 message has been received by the peer. +-- +-- This shall be called when a @PUBREC@ has been received from the peer. +-- This is the completion of the second step in a 4-way handshake. +-- The state changes from _not received_ to _not completed_. +-- We will send a @PUBREL@ to the client and expect a @PUBCOMP@ in return. +processPublishReceived :: Session auth -> PacketIdentifier -> IO () +processPublishReceived session (PacketIdentifier pid) = do + modifyMVar_ (sessionQueue session) $ \q-> + pure $! q { + notReceived = IM.delete pid (notReceived q) + , notComplete = IS.insert pid (notComplete q) + , outputBuffer = outputBuffer q Seq.|> ServerPublishRelease (PacketIdentifier pid) + } + notePending session + +-- | Release a `QoS2` message. +-- +-- This shall be called when @PUBREL@ has been received from the peer. +-- It enqueues an outgoing @PUBCOMP@. +-- The message is only released if the handler returned without exception. +-- The handler is only executed if there still is a message (it is a valid scenario +-- that it might already have been released). +processPublishRelease :: Session auth -> PacketIdentifier -> IO () +processPublishRelease session (PacketIdentifier pid) = do + modifyMVar_ (sessionQueue session) $ \q-> + case IM.lookup pid (notReleased q) of + Nothing -> + pure q + Just msg -> do + publish session msg + pure $! q { notReleased = IM.delete pid (notReleased q) + , outputBuffer = outputBuffer q Seq.|> ServerPublishComplete (PacketIdentifier pid) + } + notePending session + +-- | Complete the transmission of a QoS2 message. +-- +-- This shall be called when a @PUBCOMP@ has been received from the peer +-- to finally free the packet identifier. +processPublishComplete :: Session auth -> PacketIdentifier -> IO () +processPublishComplete session (PacketIdentifier pid) = do + modifyMVar_ (sessionQueue session) $ \q-> pure $! q { + -- The packet identifier is now free for reuse. + queuePids = PacketIdentifier pid Seq.<| queuePids q + , notComplete = IS.delete pid (notComplete q) + } + -- Although we did not enqueue something it might still be the case + -- that we have unsent data that couldn't be sent by now because no more + -- packet identifiers were available. + -- In case the output queues are actually empty, the thread will check + -- them once and immediately sleep again. + notePending session + +getSubscriptions :: Session auth -> IO (R.Trie QoS) +getSubscriptions session = + readMVar (sessionSubscriptions session) + +getConnection :: Session auth -> IO (Maybe Connection) +getConnection session = + tryReadMVar (sessionConnection session) + +getPrincipal :: Session auth -> IO Principal +getPrincipal session = + readMVar (sessionPrincipal session) + +getFreePacketIdentifiers :: Session auth -> IO (Seq.Seq PacketIdentifier) +getFreePacketIdentifiers session = + queuePids <$> readMVar (sessionQueue session) + +resetQueue :: ServerQueue -> ServerQueue +resetQueue q = q { + outputBuffer = (rePublishQoS1 . rePublishQoS2 . reReleaseQoS2) mempty + } + where + rePublishQoS1 s = IM.foldlWithKey (\s' pid msg-> s' Seq.|> ServerPublish (PacketIdentifier pid) (Duplicate True) msg) s (notAcknowledged q) + rePublishQoS2 s = IM.foldlWithKey (\s' pid msg-> s' Seq.|> ServerPublish (PacketIdentifier pid) (Duplicate True) msg) s (notReceived q) + reReleaseQoS2 s = IS.foldl (\s' pid-> s' Seq.|> ServerPublishRelease (PacketIdentifier pid) ) s (notComplete q) + +-- | This function fills the output buffer with as many messages +-- as possible (this is limited by the available packet identifiers). +normalizeQueue :: ServerQueue -> ServerQueue +normalizeQueue = takeQoS1 . takeQoS2 + where + takeQoS1 q + | Seq.null msgs = q + | otherwise = q + { outputBuffer = outputBuffer q <> Seq.zipWith (flip ServerPublish (Duplicate False)) pids' msgs' + , queuePids = pids'' + , queueQoS1 = msgs'' + , notAcknowledged = foldr (\(PacketIdentifier pid, msg)-> IM.insert pid msg) + (notAcknowledged q) + (Seq.zipWith (,) pids' msgs') + } + where + pids = queuePids q + msgs = queueQoS1 q + n = min (Seq.length pids) (Seq.length msgs) + (pids', pids'') = Seq.splitAt n pids + (msgs', msgs'') = Seq.splitAt n msgs + takeQoS2 q + | Seq.null msgs = q + | otherwise = q + { outputBuffer = outputBuffer q <> Seq.zipWith (flip ServerPublish (Duplicate False)) pids' msgs' + , queuePids = pids'' + , queueQoS2 = msgs'' + , notReceived = foldr (\(PacketIdentifier pid, msg)-> IM.insert pid msg) + (notReceived q) + (Seq.zipWith (,) pids' msgs') + } + where + pids = queuePids q + msgs = queueQoS2 q + n = min (Seq.length pids) (Seq.length msgs) + (pids', pids'') = Seq.splitAt n pids + (msgs', msgs'') = Seq.splitAt n msgs diff --git a/src/Network/MQTT/Broker/SessionStatistics.hs b/src/Network/MQTT/Broker/SessionStatistics.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Broker/SessionStatistics.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE DeriveGeneric #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Broker.SessionStatistics +-- Copyright : (c) Lars Petersen 2017 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Broker.SessionStatistics + ( Statistics () + , StatisticsSnapshot (..) + , new + , snapshot + , accountPublicationsAccepted + , accountPublicationsDropped + , accountSubscriptionsAccepted + , accountSubscriptionsDenied + , accountRetentionsAccepted + , accountRetentionsDropped ) where + +import Control.Concurrent.MVar +import qualified Data.Binary as B +import GHC.Generics + +data Statistics = Statistics + { stPublicationsAccepted :: MVar Word + , stPublicationsDropped :: MVar Word + , stSubscriptionsAccepted :: MVar Word + , stSubscriptionsDenied :: MVar Word + , stRetentionsAccepted :: MVar Word + , stRetentionsDropped :: MVar Word + } + +data StatisticsSnapshot = StatisticsSnapshot + { publicationsAccepted :: Word + , publicationsDropped :: Word + , subscriptionsAccepted :: Word + , subscriptionsDenied :: Word + , retentionsAccepted :: Word + , retentionsDropped :: Word + } deriving (Eq, Ord, Show, Generic) + +instance B.Binary StatisticsSnapshot + +new :: IO Statistics +new = Statistics + <$> newMVar 0 + <*> newMVar 0 + <*> newMVar 0 + <*> newMVar 0 + <*> newMVar 0 + <*> newMVar 0 + +snapshot :: Statistics -> IO StatisticsSnapshot +snapshot st = StatisticsSnapshot + <$> readMVar (stPublicationsAccepted st) + <*> readMVar (stPublicationsDropped st) + <*> readMVar (stSubscriptionsAccepted st) + <*> readMVar (stSubscriptionsDenied st) + <*> readMVar (stRetentionsAccepted st) + <*> readMVar (stRetentionsDropped st) + +accountPublicationsAccepted :: Statistics -> Word -> IO () +accountPublicationsAccepted ss c = + modifyMVar_ (stPublicationsAccepted ss) $ \i-> pure $! i + c + +accountPublicationsDropped :: Statistics -> Word -> IO () +accountPublicationsDropped ss c = + modifyMVar_ (stPublicationsDropped ss) $ \i-> pure $! i + c + +accountSubscriptionsAccepted :: Statistics -> Word -> IO () +accountSubscriptionsAccepted ss c = + modifyMVar_ (stSubscriptionsAccepted ss) $ \i-> pure $! i + c + +accountSubscriptionsDenied :: Statistics -> Word -> IO () +accountSubscriptionsDenied ss c = + modifyMVar_ (stSubscriptionsDenied ss) $ \i-> pure $! i + c + +accountRetentionsAccepted :: Statistics -> Word -> IO () +accountRetentionsAccepted ss c = + modifyMVar_ (stRetentionsAccepted ss) $ \i-> pure $! i + c + +accountRetentionsDropped :: Statistics -> Word -> IO () +accountRetentionsDropped ss c = + modifyMVar_ (stRetentionsDropped ss) $ \i-> pure $! i + c diff --git a/src/Network/MQTT/Client.hs b/src/Network/MQTT/Client.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Client.hs @@ -0,0 +1,445 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Client +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Client where +{- + ( Client () + , ClientConfiguration (..) + -- * Events + , ClientEvent (..) + , listenEvents + , acceptEvent + -- * Exceptions + , ClientException (..) + , newClient + , start + , stop + , publish + , subscribe + , unsubscribe + ) where + +import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.Broadcast +import Control.Exception +import Control.Monad +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BS +import qualified Data.IntMap as IM +import qualified Data.Serialize.Get as SG +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Data.Typeable +import Data.Word +import Network.MQTT +import Network.MQTT.IO +import Network.MQTT.Message +import qualified Network.Transceiver as T +import Network.URI +import System.Random + +data ClientConfiguration t + = ClientConfiguration + { -- | @mqtt://user:password\@server.example:1234@ + clientURI :: URI + , clientWill :: Maybe Message + , clientKeepAlive :: KeepAlive + , clientIdentifierPrefix :: String + , clientMaxUnacknowlegedMessages :: Word16 + , clientNewTransceiver :: IO t + } + +validateClientConfiguration :: ClientConfiguration t -> IO () +validateClientConfiguration config = do + validateUriScheme + validateUriAuthority + where + validateUriScheme = case uriScheme (clientURI config) of + "mqtt" -> pure () + "mqtts" -> pure () + "ws" -> pure () + "wss" -> pure () + _ -> err "clientURI: unsupported scheme" + validateUriAuthority = case uriAuthority (clientURI config) of + Nothing -> err "clientURI: missing authority" + Just auth -> pure () + err = throwIO . ClientConfigurationException + +data ClientConfigurationException + = ClientConfigurationException String + deriving (Typeable, Show) + +instance Exception ClientConfigurationException + +uriUsername :: URI -> Maybe Username +uriUsername uri = + f $ takeWhile (/=':') . uriUserInfo <$> uriAuthority uri + where + f (Just []) = Nothing + f (Just xs) = Just (T.pack xs) + f _ = Nothing + +uriPassword :: URI -> Maybe Password +uriPassword uri = f $ drop 1 . dropWhile (/=':') . uriUserInfo <$> uriAuthority uri + where + f (Just []) = Nothing + f (Just xs) = Just (T.encodeUtf8 $ T.pack xs) + f _ = Nothing + +data Client t + = Client + { clientIdentifier :: ClientIdentifier + , clientEventBroadcast :: Broadcast ClientEvent + , clientConfiguration :: MVar (ClientConfiguration t) + , clientRecentActivity :: MVar Bool + , clientOutput :: MVar (Either UpstreamMessage (PacketIdentifier -> (UpstreamMessage, OutboundState))) + , clientInboundState :: MVar (IM.IntMap InboundState) + , clientOutboundState :: MVar ([Int], IM.IntMap OutboundState) + , clientThreads :: MVar (Async ()) + } + +newClient :: ClientConfiguration t -> IO (Client t) +newClient configuration = Client + <$> (ClientIdentifier . T.pack . take clientIdentifierLength + . (clientIdentifierPrefix configuration ++) + . randomRs clientIdentifierCharacterRange <$> newStdGen) + <*> newBroadcast + <*> newMVar configuration + <*> newMVar False + <*> newEmptyMVar + <*> newMVar IM.empty + <*> newMVar ([0..fromIntegral (clientMaxUnacknowlegedMessages configuration)], IM.empty) + <*> (newMVar =<< async (pure ())) + where + clientIdentifierLength = 23 + clientIdentifierCharacterRange = ('a','z') + +type ClientSessionPresent = Bool + +data OutboundState + = NotAcknowledgedPublish UpstreamMessage (MVar ()) + | NotReceivedPublish UpstreamMessage (MVar ()) + | NotCompletePublish (MVar ()) + | NotAcknowledgedSubscribe UpstreamMessage (MVar [Maybe QualityOfService]) + | NotAcknowledgedUnsubscribe UpstreamMessage (MVar ()) + +newtype InboundState = NotReleasedPublish Message + +-- | Starts the `Client` if it is not already running. +-- +-- * The `Started` event will be added to the event stream. +-- * The `Connecting` event will be added to the event stream. +-- * A new `Connectable` will be created. +-- * A connection will be tried to establish (includes `connect` as well as the MQTT handshake). +-- * The `Connected` event will be added to the event stream. +-- * Messages will be processed until a `Disconnect` occurs. +-- +-- At every step after the `Connecting` event, a `Disconnect` may occur. +-- The client will then try to reconnect automatically. +start :: (T.Connectable s, T.Address s ~ a, T.StreamConnection s, T.Closable s, T.Data s ~ BS.ByteString) => Client s -> IO () +start c = modifyMVar_ (clientThreads c) $ \p-> + poll p >>= \m-> case m of + -- Processing thread is stil running, no need to connect. + Nothing -> pure p + -- Processing thread not running, create a new one with new connection + Just _ -> do + broadcast (clientEventBroadcast c) Started + async $ forever $ do + run -- `catch` (\e-> print (e :: SomeException) >> print "RECONNECT") + threadDelay 1000000 + + where + run :: IO () + run = join (clientNewTransceiver <$> readMVar (clientConfiguration c)) >>= handleConnection False + + -- handleConnection :: (StreamTransmitter s, StreamReceiver s, Connectable s, Closable s) => ClientSessionPresent -> s -> IO () + handleConnection clientSessionPresent connection = do + broadcast (clientEventBroadcast c) Connecting + connectTransmitter + broadcast (clientEventBroadcast c) Connected + sendConnect + receiveConnectAcknowledgement >>= maintainConnection + where + connectTransmitter :: IO () + connectTransmitter = + T.connect connection =<< readMVar undefined + + sendConnect :: IO () + sendConnect = do + conf <- readMVar (clientConfiguration c) + T.sendChunks connection $ BS.toLazyByteString $ buildUpstreamMessage Connect + { connectClientIdentifier = clientIdentifier c + , connectCleanSession = False + , connectKeepAlive = clientKeepAlive conf + , connectWill = clientWill conf + , connectCredentials = undefined + } + + receiveConnectAcknowledgement :: IO BS.ByteString + receiveConnectAcknowledgement = do + bs <- T.receiveChunk connection + case SG.runGetPartial parseDownstreamMessage bs of + SG.Done message bs' -> f message >> pure bs' + SG.Fail e bs' -> throwIO $ ProtocolViolation e + SG.Partial _ -> throwIO $ ProtocolViolation "Expected CONNACK, got end of input." + where + f (ConnectAcknowledgement a) = case a of + Right serverSessionPresent + | serverSessionPresent && not clientSessionPresent -> + -- This should not happen unless an old client identifier + -- has been chosen or the server is wicked. + throwIO ClientLostSession + | not serverSessionPresent && clientSessionPresent -> + -- This may happen if the server lost its memory, either + -- acidentially or for administrative reasons. + throwIO ServerLostSession + | otherwise -> pure () + Left connectionRefusal -> throwIO $ ConnectionRefused connectionRefusal + f _ = throwIO $ ProtocolViolation "Expected CONNACK, got something else." + + maintainConnection :: BS.ByteString -> IO () + maintainConnection i = + keepAlive `race_` handleOutput `race_` (handleInput i `catch` \e-> print (e :: SomeException)) + where + -- The keep alive thread wakes up every `keepAlive/2` seconds. + -- It reads and unsets the recent-activity flag. + -- When it finds the recent-activity flag unset, it sends a PINGREQ to the + -- server. The interval between the last message and the PINGREQ is at + -- most `keepAlive` seconds (assuming we get woken up on time by the RTS). + keepAlive :: IO () + keepAlive = do + interval <- (500000*) . fromIntegral . clientKeepAlive <$> readMVar (clientConfiguration c) + forever $ do + threadDelay interval + activity <- swapMVar (clientRecentActivity c) False + unless activity $ putMVar (clientOutput c) $ Left PingRequest + + handleOutput :: IO () + handleOutput = bufferedOutput connection getMessage getMaybeMessage (T.sendChunk connection) + where + getMessage :: IO UpstreamMessage + getMessage = do + msg <- f =<< takeMVar (clientOutput c) + void $ swapMVar (clientRecentActivity c) True + pure msg + + getMaybeMessage :: IO (Maybe UpstreamMessage) + getMaybeMessage = do + memsg <- tryTakeMVar (clientOutput c) + case memsg of + Nothing -> pure Nothing + Just emsg -> Just <$> f emsg + + f :: Either UpstreamMessage (PacketIdentifier -> (UpstreamMessage, OutboundState)) -> IO UpstreamMessage + f (Left msg) = pure msg + f (Right imsg) = assignPacketIdentifier imsg + + assignPacketIdentifier :: (PacketIdentifier -> (UpstreamMessage, OutboundState)) -> IO UpstreamMessage + assignPacketIdentifier x = + modifyMVar (clientOutboundState c) assign >>= \mm-> case mm of + Just m -> pure m + -- We cannot easily wait for when packet identifiers are available again. + -- On the other hand throwing an exception seems too drastic. So (for the + -- extremely unlikely) case of packet identifier exhaustion, we shall wait + -- 100ms and then try again. + Nothing -> threadDelay 100000 >> assignPacketIdentifier x + where + assign im@([], _) = pure (im, Nothing) + assign (i:is, m) = let (msg, st) = x (PacketIdentifier i) + in pure ((is, IM.insert i st m), Just msg) + + handleInput :: BS.ByteString -> IO () + handleInput bs + | BS.null bs = handleInput' =<< T.receiveChunk connection + | otherwise = handleInput' bs + where + handleInput' bs' = do + -- print $ "handle input" ++ show (BS.unpack bs') + g (SG.runGetPartial parseDownstreamMessage bs') + g (SG.Done message bs') = do + -- print message + f message >> handleInput' bs' + g (SG.Partial cont) = do + --print "Partial" + bs' <- T.receiveChunk connection + if BS.null bs' + then throwIO ServerClosedConnection + else g $ cont bs' + g (SG.Fail e _) = do + --print $ "FAIL" ++ show e + throwIO $ ProtocolViolation e + + f msg@Publish {} = case publishQoS msg of + PublishQoS0 -> broadcast (clientEventBroadcast c) $ Received Message + { topic = publishTopic msg + , payload = publishBody msg + , qos = QoS0 + , retained = publishRetain msg + , duplicate = False + } + {-PublishQoS1 dup i -> do + broadcast (clientEventBroadcast c) $ Received Message + { topic = publishTopic msg + , payload = publishBody msg + , qos = QoS1 + , retained = publishRetain msg + , duplicate = False + } + putMVar (clientOutput c) $ Left $ PublishAcknowledgement i + PublishQoS2 (PacketIdentifier i) -> + modifyMVar_ (clientInboundState c) $ + pure . IM.insert i (NotReleasedPublish Message + { topic = publishTopic msg + , payload = publishBody msg + , qos = QoS2 + , retained = publishRetain msg + , duplicate = False + }) + -- The following packet types are responses to earlier requests. + -- We need to dispatch them to the waiting threads. + f (PublishAcknowledgement (PacketIdentifier i)) = + modifyMVar_ (clientOutboundState c) $ \(is,im)-> + case IM.lookup i im of + Just (NotAcknowledgedPublish _ promise) -> + putMVar promise () >> pure (i:is, IM.delete i im) + _ -> throwIO $ ProtocolViolation "Expected PUBACK, got something else." + f (PublishReceived (PacketIdentifier i)) = do + modifyMVar_ (clientOutboundState c) $ \(is,im)-> + case IM.lookup i im of + Just (NotReceivedPublish _ promise) -> + pure (is, IM.insert i (NotCompletePublish promise) im) + _ -> throwIO $ ProtocolViolation "Expected PUBREC, got something else." + putMVar (clientOutput c) (Left $ PublishRelease (PacketIdentifier i)) + f (PublishRelease (PacketIdentifier i)) = do + modifyMVar_ (clientInboundState c) $ \im-> + case IM.lookup i im of + Just (NotReleasedPublish msg) -> do + broadcast (clientEventBroadcast c) $ Received msg -- Publish exactly once here! + pure (IM.delete i im) + Nothing -> -- Duplicate, don't publish again. + pure im + putMVar (clientOutput c) (Left $ PublishComplete (PacketIdentifier i)) -} + f (PublishComplete (PacketIdentifier i)) = do + modifyMVar_ (clientOutboundState c) $ \p@(is,im)-> + case IM.lookup i im of + Nothing -> + pure p + Just (NotCompletePublish future) -> do + putMVar future () + pure (i:is, IM.delete i im) + _ -> + throwIO $ ProtocolViolation "Expected PUBCOMP, got something else." + f (SubscribeAcknowledgement (PacketIdentifier i) as) = + modifyMVar_ (clientOutboundState c) $ \p@(is,im)-> + case IM.lookup i im of + Nothing -> pure p + Just (NotAcknowledgedSubscribe m promise) -> do + putMVar promise as + pure (i:is, IM.delete i im) + _ -> throwIO $ ProtocolViolation "Expected PUBCOMP, got something else." + f (UnsubscribeAcknowledgement (PacketIdentifier i)) = + modifyMVar_ (clientOutboundState c) $ \p@(is,im)-> + case IM.lookup i im of + Nothing -> pure p + Just (NotAcknowledgedUnsubscribe m promise) -> do + putMVar promise () + pure (i:is, IM.delete i im) + _ -> throwIO $ ProtocolViolation "Expected PUBCOMP, got something else." + f PingResponse = pure () + -- The following packets must not be sent from the server to the client. + _ = throwIO $ ProtocolViolation "Unexpected packet type received from the server." + +-- | Stops the `Client` if it is not already stopped. +-- +-- * A graceful DISCONNECT packet will be sent to the server and the +-- connection will be `close`d. +-- * All internal client threads will be terminated. +-- * The `Stopped` event will be added to the event stream. +stop :: T.Closable s => Client s -> IO () +stop c = do + t <- readMVar (clientThreads c) + putMVar (clientOutput c) (Left Disconnect) + wait t + +publish :: Client s -> Topic -> QoS -> Retain -> Payload -> IO () +publish client !topic !qos !retain !payload = case qos of + QoS0 -> undefined -- putMVar (clientOutput client) $ Left $ message Nothing + QoS1 -> undefined {-- register AtLeastOnce NotAcknowledgedPublish + QoS2 -> register ExactlyOnce NotReceivedPublish + where + register qos' f = do + confirmation <- newEmptyMVar + putMVar (clientOutput client) $ Right $ qosMessage confirmation qos' f + takeMVar confirmation + qosMessage confirmation qos f packetIdentifier = + ( m, f m { publishDuplicate = True } confirmation ) + where + m = message $ Just (qos, packetIdentifier) + message mqos' = Publish { + publishDuplicate = False, + publishRetain = retain, + publishQoS = mqos', + publishTopic = topic, + publishBody = body } -} + +subscribe :: Client s -> [(TopicFilter, QoS)] -> IO [Maybe QoS] +subscribe client [] = pure [] +subscribe client topics = do + response <- newEmptyMVar + putMVar (clientOutput client) $ Right $ f response + takeMVar response + where + f response i = + let message = Subscribe i topics + in (message, NotAcknowledgedSubscribe message response) + +unsubscribe :: Client s -> [TopicFilter] -> IO () +unsubscribe client [] = pure () +unsubscribe client topics = do + confirmation <- newEmptyMVar + putMVar (clientOutput client) $ Right $ f confirmation + takeMVar confirmation + where + f confirmation i = + let message = Unsubscribe i topics + in (message, NotAcknowledgedUnsubscribe message confirmation) + +listenEvents :: Client s -> IO (BroadcastListener ClientEvent) +listenEvents c = listen (clientEventBroadcast c) + +acceptEvent :: BroadcastListener ClientEvent -> IO ClientEvent +acceptEvent = accept + +data ClientEvent + = Started + | Connecting + | Connected + | Ping + | Pong + | Received Message + | Disconnected ClientException + | Stopped + deriving (Eq, Ord, Show) + +data ClientException + = ProtocolViolation String + | ConnectionRefused ConnectionRefusal + | ClientLostSession + | ServerLostSession + | ClientClosedConnection + | ServerClosedConnection + deriving (Eq, Ord, Show, Typeable) + +instance Exception ClientException where +-} diff --git a/src/Network/MQTT/Message.hs b/src/Network/MQTT/Message.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Message.hs @@ -0,0 +1,549 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TupleSections #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Message +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Message ( + -- * Message + Message (..) + , module Network.MQTT.Message.Topic + -- ** QoS + , QoS (..) + -- ** Retain + , Retain (..) + -- ** Payload + , Payload (..) + + -- * Packet + -- ** ClientPacket + , ClientPacket (..) + , clientPacketBuilder + , clientPacketParser + -- ** ServerPacket + , ServerPacket (..) + , serverPacketBuilder + , serverPacketParser + + -- ** PacketIdentifier + , PacketIdentifier (..) + -- ** ClientIdentifier + , ClientIdentifier (..) + -- ** Username + , Username (..) + -- ** Password + , Password (..) + -- ** CleanSession + , CleanSession (..) + -- ** SessionPresent + , SessionPresent (..) + -- ** RejectReason + , RejectReason (..) + -- ** Duplicate + , Duplicate (..) + -- ** KeepAliveInterval + , KeepAliveInterval (..) + + -- * Other (internal) exports + , lengthParser + , lengthBuilder + , utf8Parser + , utf8Builder + ) where + +import Control.Monad +import qualified Data.Attoparsec.ByteString as A +import qualified Data.Binary as B +import qualified Data.Binary.Get as SG +import Data.Bits +import Data.Bool +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BS +import qualified Data.ByteString.Lazy as BSL +import Data.Monoid +import Data.String +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Data.Word +import GHC.Generics + +import Network.MQTT.Message.QoS +import Network.MQTT.Message.Topic +import qualified Network.MQTT.Message.Topic as TF + +newtype SessionPresent = SessionPresent Bool deriving (Eq, Ord, Show) +newtype CleanSession = CleanSession Bool deriving (Eq, Ord, Show) +newtype Retain = Retain Bool deriving (Eq, Ord, Show) +newtype Payload = Payload BSL.ByteString deriving (Eq, Ord, Show, IsString) +newtype Duplicate = Duplicate Bool deriving (Eq, Ord, Show) +newtype KeepAliveInterval = KeepAliveInterval Word16 deriving (Eq, Ord, Show, Num) +newtype Username = Username T.Text deriving (Eq, Ord, Show, IsString, Generic) +newtype Password = Password BS.ByteString deriving (Eq) +newtype ClientIdentifier = ClientIdentifier T.Text deriving (Eq, Ord, Show, IsString, Generic) +newtype PacketIdentifier = PacketIdentifier Int deriving (Eq, Ord, Show) + +instance Show Password where + show = const "*********" + +instance B.Binary ClientIdentifier +instance B.Binary Username + +data RejectReason + = UnacceptableProtocolVersion + | IdentifierRejected + | ServerUnavailable + | BadUsernameOrPassword + | NotAuthorized + deriving (Eq, Ord, Show, Enum, Bounded) + +data Message + = Message + { msgTopic :: !TF.Topic + , msgQoS :: !QoS + , msgRetain :: !Retain + , msgPayload :: !Payload + } deriving (Eq, Ord, Show) + +data ClientPacket + = ClientConnect + { connectClientIdentifier :: !ClientIdentifier + , connectCleanSession :: !CleanSession + , connectKeepAlive :: !KeepAliveInterval + , connectWill :: !(Maybe Message) + , connectCredentials :: !(Maybe (Username, Maybe Password)) + } + | ClientConnectUnsupported + | ClientPublish {-# UNPACK #-} !PacketIdentifier !Duplicate !Message + | ClientPublishAcknowledged {-# UNPACK #-} !PacketIdentifier + | ClientPublishReceived {-# UNPACK #-} !PacketIdentifier + | ClientPublishRelease {-# UNPACK #-} !PacketIdentifier + | ClientPublishComplete {-# UNPACK #-} !PacketIdentifier + | ClientSubscribe {-# UNPACK #-} !PacketIdentifier ![(TF.Filter, QoS)] + | ClientUnsubscribe {-# UNPACK #-} !PacketIdentifier ![TF.Filter] + | ClientPingRequest + | ClientDisconnect + deriving (Eq, Show) + +data ServerPacket + = ServerConnectionAccepted !SessionPresent + | ServerConnectionRejected !RejectReason + | ServerPublish {-# UNPACK #-} !PacketIdentifier !Duplicate !Message + | ServerPublishAcknowledged {-# UNPACK #-} !PacketIdentifier + | ServerPublishReceived {-# UNPACK #-} !PacketIdentifier + | ServerPublishRelease {-# UNPACK #-} !PacketIdentifier + | ServerPublishComplete {-# UNPACK #-} !PacketIdentifier + | ServerSubscribeAcknowledged {-# UNPACK #-} !PacketIdentifier ![Maybe QoS] + | ServerUnsubscribeAcknowledged {-# UNPACK #-} !PacketIdentifier + | ServerPingResponse + deriving (Eq, Show) + +clientPacketParser :: SG.Get ClientPacket +clientPacketParser = + SG.lookAhead SG.getWord8 >>= \h-> case h .&. 0xf0 of + 0x10 -> connectParser + 0x30 -> publishParser ClientPublish + 0x40 -> acknowledgedParser ClientPublishAcknowledged + 0x50 -> acknowledgedParser ClientPublishReceived + 0x60 -> acknowledgedParser ClientPublishRelease + 0x70 -> acknowledgedParser ClientPublishComplete + 0x80 -> subscribeParser + 0xa0 -> unsubscribeParser + 0xc0 -> pingRequestParser + 0xe0 -> disconnectParser + _ -> fail "clientPacketParser: Invalid message type." + +serverPacketParser :: SG.Get ServerPacket +serverPacketParser = + SG.lookAhead SG.getWord8 >>= \h-> case h .&. 0xf0 of + 0x20 -> connectAcknowledgedParser + 0x30 -> publishParser ServerPublish + 0x40 -> acknowledgedParser ServerPublishAcknowledged + 0x50 -> acknowledgedParser ServerPublishReceived + 0x60 -> acknowledgedParser ServerPublishRelease + 0x70 -> acknowledgedParser ServerPublishComplete + 0xb0 -> acknowledgedParser ServerUnsubscribeAcknowledged + 0x90 -> subscribeAcknowledgedParser + 0xd0 -> pingResponseParser + _ -> fail "serverPacketParser: Packet type not implemented." + +connectParser :: SG.Get ClientPacket +connectParser = do + h <- SG.getWord8 + when (h .&. 0x0f /= 0) $ + fail "connectParser: The header flags are reserved and MUST be set to 0." + void lengthParser -- the remaining length is redundant in this packet type + y <- SG.getWord64be -- get the next 8 bytes all at once and not byte per byte + case y .&. 0xffffffffffffff00 of + 0x00044d5154540400 -> do + let cleanSession = CleanSession ( y .&. 0x02 /= 0 ) + let retain = Retain ( y .&. 0x20 /= 0 ) + keepAlive <- KeepAliveInterval <$> SG.getWord16be + cid <- ClientIdentifier <$> utf8Parser + will <- if y .&. 0x04 == 0 + then pure Nothing + else Just <$> do + topic <- fst <$> topicAndLengthParser + bodyLen <- fromIntegral <$> SG.getWord16be + payload <- Payload <$> SG.getLazyByteString bodyLen + qos <- case y .&. 0x18 of + 0x00 -> pure QoS0 + 0x08 -> pure QoS1 + 0x10 -> pure QoS2 + _ -> fail "connectParser: Violation of [MQTT-3.1.2-14]." + pure $ Message topic qos retain payload + cred <- if y .&. 0x80 == 0 + then pure Nothing + else Just <$> ( (,) + <$> (Username <$> utf8Parser) + <*> if y .&. 0x40 == 0 + then pure Nothing + else Just . Password <$> (SG.getByteString . fromIntegral =<< SG.getWord16be) ) + pure ( ClientConnect cid cleanSession keepAlive will cred ) + -- This is the prefix of the version 3 protocol. + -- We just assume that this is version 3 and return immediately. + -- The caller shall send return code 0x01 (unacceptable protocol). + 0x00064d5149736400 -> pure ClientConnectUnsupported + -- This case is different from the previous as it is either not + -- MQTT or a newer protocol version we don't know (yet). + -- The caller shall close the connection immediately without + -- sending any data in this case. + _ -> fail "connectParser: Unexpected protocol initialization." + +connectAcknowledgedParser :: SG.Get ServerPacket +connectAcknowledgedParser = do + x <- SG.getWord32be + case x .&. 0xff of + 0 -> pure $ ServerConnectionAccepted $ SessionPresent (x .&. 0x0100 /= 0) + 1 -> pure $ ServerConnectionRejected UnacceptableProtocolVersion + 2 -> pure $ ServerConnectionRejected IdentifierRejected + 3 -> pure $ ServerConnectionRejected ServerUnavailable + 4 -> pure $ ServerConnectionRejected BadUsernameOrPassword + 5 -> pure $ ServerConnectionRejected NotAuthorized + _ -> fail "connectAcknowledgedParser: Invalid (reserved) return code." + +publishParser :: (PacketIdentifier -> Duplicate -> Message -> a) -> SG.Get a +publishParser publish = do + hflags <- SG.getWord8 + let dup = Duplicate $ hflags .&. 0x08 /= 0 -- duplicate flag + let ret = Retain $ hflags .&. 0x01 /= 0 -- retain flag + len <- lengthParser + (topic, topicLen) <- topicAndLengthParser + let qosBits = hflags .&. 0x06 + if qosBits == 0x00 + then do + body <- SG.getLazyByteString $ fromIntegral ( len - topicLen ) + pure $ publish (PacketIdentifier (-1)) dup Message { + msgTopic = topic + , msgPayload = Payload body + , msgQoS = QoS0 + , msgRetain = ret + } + else do + let qos = if qosBits == 0x02 then QoS1 else QoS2 + pid <- PacketIdentifier . fromIntegral <$> SG.getWord16be + body <- SG.getLazyByteString $ fromIntegral ( len - topicLen - 2 ) + pure $ publish pid dup Message { + msgTopic = topic + , msgPayload = Payload body + , msgQoS = qos + , msgRetain = ret + } + +acknowledgedParser :: (PacketIdentifier -> a) -> SG.Get a +acknowledgedParser f = do + w32 <- SG.getWord32be + pure $ f $ PacketIdentifier $ fromIntegral $ w32 .&. 0xffff +{-# INLINE acknowledgedParser #-} + +subscribeParser :: SG.Get ClientPacket +subscribeParser = do + void SG.getWord8 + rlen <- lengthParser + pid <- PacketIdentifier . fromIntegral <$> SG.getWord16be + ClientSubscribe pid <$> parseFilters (rlen - 2) [] + where + parseFilters r accum + | r <= 0 = pure (reverse accum) + | otherwise = do + (filtr,len) <- filterAndLengthParser + qos <- getQoS + parseFilters ( r - 1 - len ) ( ( filtr, qos ) : accum ) + getQoS = SG.getWord8 >>= \w-> case w of + 0x00 -> pure QoS0 + 0x01 -> pure QoS1 + 0x02 -> pure QoS2 + _ -> fail "clientSubscribeParser: Violation of [MQTT-3.8.3-4]." + +unsubscribeParser :: SG.Get ClientPacket +unsubscribeParser = do + void SG.getWord8 + rlen <- lengthParser + pid <- PacketIdentifier . fromIntegral <$> SG.getWord16be + ClientUnsubscribe pid <$> parseFilters (rlen - 2) [] + where + parseFilters r accum + | r <= 0 = pure (reverse accum) + | otherwise = do + (filtr,len) <- filterAndLengthParser + parseFilters ( r - len ) ( filtr : accum ) + +subscribeAcknowledgedParser :: SG.Get ServerPacket +subscribeAcknowledgedParser = do + void SG.getWord8 + rlen <- lengthParser + pid <- PacketIdentifier . fromIntegral <$> SG.getWord16be + ServerSubscribeAcknowledged pid <$> (map f . BS.unpack <$> SG.getByteString (rlen - 2)) + where + f 0x00 = Just QoS0 + f 0x01 = Just QoS1 + f 0x02 = Just QoS2 + f _ = Nothing + +pingRequestParser :: SG.Get ClientPacket +pingRequestParser = do + void SG.getWord16be + pure ClientPingRequest +{-# INLINE pingRequestParser #-} + +pingResponseParser :: SG.Get ServerPacket +pingResponseParser = do + void SG.getWord16be + pure ServerPingResponse +{-# INLINE pingResponseParser #-} + +disconnectParser :: SG.Get ClientPacket +disconnectParser = do + void SG.getWord16be + pure ClientDisconnect +{-# INLINE disconnectParser #-} + +clientPacketBuilder :: ClientPacket -> BS.Builder +clientPacketBuilder (ClientConnect + (ClientIdentifier cid) + (CleanSession cleanSession) + (KeepAliveInterval keepAlive) will credentials) = + BS.word8 0x10 + <> lengthBuilder ( 10 + cidLen + willLen + credLen ) + <> BS.word64BE ( 0x00044d5154540400 .|. willFlag .|. credFlag .|. sessFlag ) + <> BS.word16BE keepAlive + <> cidBuilder + <> willBuilder + <> credBuilder + where + sessFlag = if cleanSession then 0x02 else 0x00 + (cidBuilder, cidLen) = utf8Builder cid + (willBuilder, willLen, willFlag) = case will of + Nothing -> + (mempty, 0, 0x00) + Just (Message t q (Retain r) (Payload b))-> + let tlen = TF.topicLength t + blen = fromIntegral (BSL.length b) + qflag = case q of + QoS0 -> 0x04 + QoS1 -> 0x0c + QoS2 -> 0x14 + rflag = if r then 0x20 else 0x00 + x1 = BS.word16BE (fromIntegral tlen) + x2 = TF.topicBuilder t + x3 = BS.word16BE (fromIntegral blen) + x4 = BS.lazyByteString b + in (x1 <> x2 <> x3 <> x4, 4 + tlen + blen, qflag .|. rflag) + (credBuilder, credLen, credFlag) = case credentials of + Nothing -> + (mempty, 0, 0x00) + Just (Username ut, Nothing) -> + let u = T.encodeUtf8 ut + ulen = BS.length u + x1 = BS.word16BE (fromIntegral ulen) + x2 = BS.byteString u + in (x1 <> x2, 2 + ulen, 0x80) + Just (Username ut, Just (Password p)) -> + let u = T.encodeUtf8 ut + ulen = BS.length u + plen = BS.length p + x1 = BS.word16BE (fromIntegral ulen) + x2 = BS.byteString u + x3 = BS.word16BE (fromIntegral plen) + x4 = BS.byteString p + in (x1 <> x2 <> x3 <> x4, 4 + ulen + plen, 0xc0) +clientPacketBuilder ClientConnectUnsupported = + BS.word8 10 <> lengthBuilder 38 <> BS.word64BE 0x00064d514973647063 +clientPacketBuilder (ClientPublish pid dup msg) = + publishBuilder pid dup msg +clientPacketBuilder (ClientPublishAcknowledged (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x40020000 .|. pid +clientPacketBuilder (ClientPublishReceived (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x50020000 .|. pid +clientPacketBuilder (ClientPublishRelease (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x62020000 .|. pid +clientPacketBuilder (ClientPublishComplete (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x70020000 .|. pid +clientPacketBuilder (ClientSubscribe (PacketIdentifier pid) filters) = + BS.word8 0x82 <> lengthBuilder len <> BS.word16BE (fromIntegral pid) + <> mconcat (fmap filterBuilder' filters) + where + filterBuilder' (f, q) = BS.word16BE (fromIntegral fl) <> fb <> qb + where + fb = TF.filterBuilder f + fl = TF.filterLength f + qb = BS.word8 $ case q of + QoS0 -> 0x00 + QoS1 -> 0x01 + QoS2 -> 0x02 + len = 2 + sum ( map ( (+3) . TF.filterLength . fst ) filters ) +clientPacketBuilder (ClientUnsubscribe (PacketIdentifier pid) filters) = + BS.word8 0xa2 <> lengthBuilder len <> BS.word16BE (fromIntegral pid) + <> mconcat ( map filterBuilder' filters ) + where + filterBuilder' f = BS.word16BE (fromIntegral fl) <> fb + where + fb = TF.filterBuilder f + fl = TF.filterLength f + len = 2 + sum ( map ( (+2) . TF.filterLength ) filters ) +clientPacketBuilder ClientPingRequest = + BS.word16BE 0xc000 +clientPacketBuilder ClientDisconnect = + BS.word16BE 0xe000 + +serverPacketBuilder :: ServerPacket -> BS.Builder +serverPacketBuilder (ServerConnectionAccepted (SessionPresent sessionPresent)) + | sessionPresent = BS.word32BE 0x20020100 + | otherwise = BS.word32BE 0x20020000 +serverPacketBuilder (ServerConnectionRejected reason) = + BS.word32BE $ case reason of + UnacceptableProtocolVersion -> 0x20020001 + IdentifierRejected -> 0x20020002 + ServerUnavailable -> 0x20020003 + BadUsernameOrPassword -> 0x20020004 + NotAuthorized -> 0x20020005 +serverPacketBuilder (ServerPublish pid dup msg) = + publishBuilder pid dup msg +serverPacketBuilder (ServerPublishAcknowledged (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x40020000 .|. pid +serverPacketBuilder (ServerPublishReceived (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x50020000 .|. pid +serverPacketBuilder (ServerPublishRelease (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x62020000 .|. pid +serverPacketBuilder (ServerPublishComplete (PacketIdentifier pid)) = + BS.word32BE $ fromIntegral $ 0x70020000 .|. pid +serverPacketBuilder (ServerSubscribeAcknowledged (PacketIdentifier pid) rcs) = + BS.word8 0x90 <> lengthBuilder (2 + length rcs) + <> BS.word16BE (fromIntegral pid) <> mconcat ( map ( BS.word8 . f ) rcs ) + where + f Nothing = 0x80 + f (Just QoS0) = 0x00 + f (Just QoS1) = 0x01 + f (Just QoS2) = 0x02 +serverPacketBuilder (ServerUnsubscribeAcknowledged (PacketIdentifier pid)) = + BS.word16BE 0xb002 <> BS.word16BE (fromIntegral pid) +serverPacketBuilder ServerPingResponse = + BS.word16BE 0xd000 + +publishBuilder :: PacketIdentifier -> Duplicate -> Message -> BS.Builder +publishBuilder (PacketIdentifier pid) (Duplicate dup) msg = + BS.word8 h + <> lengthBuilder len + <> BS.word16BE (fromIntegral topicLen) + <> topicBuilder' + <> bool (BS.word16BE $ fromIntegral pid) mempty (msgQoS msg == QoS0) + <> BS.lazyByteString body + where + Payload body = msgPayload msg + topicLen = TF.topicLength (msgTopic msg) + topicBuilder' = TF.topicBuilder (msgTopic msg) + len = 2 + topicLen + fromIntegral (BSL.length body) + + bool 2 0 (msgQoS msg == QoS0) + h = 0x30 + .|. bool 0x00 0x01 retain + .|. bool 0x00 0x08 dup + .|. case msgQoS msg of + QoS0 -> 0x00 + QoS1 -> 0x02 + QoS2 -> 0x04 + where + Retain retain = msgRetain msg + +topicAndLengthParser :: SG.Get (TF.Topic, Int) +topicAndLengthParser = do + topicLen <- fromIntegral <$> SG.getWord16be + topicBytes <- SG.getByteString topicLen + case A.parseOnly TF.topicParser topicBytes of + Right t -> pure (t, topicLen + 2) + Left _ -> fail "topicAndLengthParser: Invalid topic." +{-# INLINE topicAndLengthParser #-} + +filterAndLengthParser :: SG.Get (TF.Filter, Int) +filterAndLengthParser = do + filterLen <- fromIntegral <$> SG.getWord16be + filterBytes <- SG.getByteString filterLen + case A.parseOnly TF.filterParser filterBytes of + Right f -> pure (f, filterLen + 2) + Left _ -> fail "filterAndLengthParser: Invalid filter." +{-# INLINE filterAndLengthParser #-} + +utf8Parser :: SG.Get T.Text +utf8Parser = do + str <- SG.getWord16be >>= SG.getByteString . fromIntegral + when (BS.elem 0x00 str) (fail "utf8Parser: Violation of [MQTT-1.5.3-2].") + case T.decodeUtf8' str of + Right txt -> return txt + _ -> fail "utf8Parser: Violation of [MQTT-1.5.3]." +{-# INLINE utf8Parser #-} + +utf8Builder :: T.Text -> (BS.Builder, Int) +utf8Builder txt = + if len > 0xffff + then error "utf8Builder: Encoded size must be <= 0xffff." + else (BS.word16BE (fromIntegral len) <> BS.byteString bs, len + 2) + where + bs = T.encodeUtf8 txt + len = BS.length bs +{-# INLINE utf8Builder #-} + +lengthParser:: SG.Get Int +lengthParser = do + b0 <- fromIntegral <$> SG.getWord8 + if b0 < 128 + then pure b0 + else do + b1 <- fromIntegral <$> SG.getWord8 + if b1 < 128 + then pure $ b1 * 128 + (b0 .&. 127) + else do + b2 <- fromIntegral <$> SG.getWord8 + if b2 < 128 + then pure $ b2 * 128 * 128 + (b1 .&. 127) * 128 + (b0 .&. 127) + else do + b3 <- fromIntegral <$> SG.getWord8 + if b3 < 128 + then pure $ b3 * 128 * 128 * 128 + (b2 .&. 127) * 128 * 128 + (b1 .&. 127) * 128 + (b0 .&. 127) + else fail "lengthParser: invalid input" +{-# INLINE lengthParser #-} + +lengthBuilder :: Int -> BS.Builder +lengthBuilder i + | i < 0x80 = BS.word8 ( fromIntegral i ) + | i < 0x80*0x80 = BS.word16LE $ fromIntegral $ 0x0080 -- continuation bit + .|. ( i .&. 0x7f ) + .|. unsafeShiftL ( i .&. 0x3f80 ) 1 + | i < 0x80*0x80*0x80 = BS.word16LE ( fromIntegral $ 0x8080 + .|. ( i .&. 0x7f ) + .|. unsafeShiftL ( i .&. 0x3f80 ) 1 + ) + <> BS.word8 ( fromIntegral + $ unsafeShiftR ( i .&. 0x1fc000 ) 14 + ) + | i < 0x80*0x80*0x80*0x80 = BS.word32LE $ fromIntegral $ 0x00808080 + .|. ( i .&. 0x7f ) + .|. unsafeShiftL ( i .&. 0x3f80 ) 1 + .|. unsafeShiftL ( i .&. 0x1fc000 ) 2 + .|. unsafeShiftL ( i .&. 0x0ff00000) 3 + | otherwise = error "lengthBuilder: invalid input" +{-# INLINE lengthBuilder #-} diff --git a/src/Network/MQTT/Message/QoS.hs b/src/Network/MQTT/Message/QoS.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Message/QoS.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE TypeFamilies #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Message +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Message.QoS where + +import Network.MQTT.Trie + +-- | The quality of service defines the guarantees given wrt to message reception. +data QoS + = QoS0 -- ^ Message delivery is not guaranteed. + | QoS1 -- ^ Message is guaranteed to be delivered at least once. + | QoS2 -- ^ Message is guaranteed to be delivered exactly once. + deriving (Eq, Ord, Show, Enum, Bounded) + +instance TrieValue QoS where + data TrieNode QoS = QosNode {-# UNPACK #-} !Int !(Trie QoS) + node t (Just QoS0) = QosNode 0 t + node t (Just QoS1) = QosNode 1 t + node t (Just QoS2) = QosNode 2 t + node t Nothing = QosNode 3 t + nodeNull = const False + nodeTree (QosNode _ t) = t + nodeValue (QosNode 0 _) = Just QoS0 + nodeValue (QosNode 1 _) = Just QoS1 + nodeValue (QosNode 2 _) = Just QoS2 + nodeValue (QosNode _ _) = Nothing diff --git a/src/Network/MQTT/Message/Topic.hs b/src/Network/MQTT/Message/Topic.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Message/Topic.hs @@ -0,0 +1,181 @@ +{-# LANGUAGE OverloadedStrings #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Topic +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Message.Topic + ( -- ** Topic + Topic () + , topicLevels + , topicLength + , topicParser + , topicBuilder + -- ** Filter + , Filter () + , filterLevels + , filterLength + , filterParser + , filterBuilder + -- ** Level + , Level () + , levelParser + , multiLevelWildcard + , singleLevelWildcard + , startsWithDollar + ) where + +import Data.Monoid ((<>)) +import Control.Applicative +import Control.Monad (void) +import qualified Data.Attoparsec.ByteString as A +import qualified Data.ByteString.Builder as BS +import qualified Data.ByteString.Short as BS +import Data.List +import Data.List.NonEmpty (NonEmpty (..)) +import Data.String +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.Encoding.Error as T +import Data.Word + +-- | According to the MQTT specification a topic +-- +-- * must not be empty +-- * must not contain @+@, @#@ or @\\NUL@ characters +newtype Topic = Topic (NonEmpty Level) deriving (Eq, Ord) +-- | According to the MQTT specification a filter +-- +-- * must not be empty +-- * must not contain a @\\NUL@ character +newtype Filter = Filter (NonEmpty Level) deriving (Eq, Ord) +-- | A `Level` represents a single path element of a `Topic` or a `Filter`. +newtype Level = Level BS.ShortByteString deriving (Eq, Ord) + +instance Show Topic where + show (Topic xs) = show (Filter xs) + +instance Show Filter where + show (Filter (x:|xs)) = concat ["\"", intercalate "/" $ f x : map f xs, "\""] + where + f (Level l) = T.unpack $ T.decodeUtf8With T.lenientDecode $ BS.fromShort l + +instance Show Level where + show (Level x) = + concat ["\"", T.unpack $ T.decodeUtf8With T.lenientDecode $ BS.fromShort x, "\""] + +instance IsString Topic where + fromString s = case A.parseOnly topicParser (T.encodeUtf8 $ T.pack s) of + Left e -> error e + Right t -> t + +instance IsString Filter where + fromString s = case A.parseOnly filterParser (T.encodeUtf8 $ T.pack s) of + Left e -> error e + Right t -> t + +instance IsString Level where + fromString s = case A.parseOnly levelParser (T.encodeUtf8 $ T.pack s) of + Left e -> error e + Right t -> t + +topicLevels :: Topic -> NonEmpty Level +topicLevels (Topic x) = x +{-# INLINE topicLevels #-} + +filterLevels :: Filter -> NonEmpty Level +filterLevels (Filter x) = x +{-# INLINE filterLevels #-} + +topicParser :: A.Parser Topic +topicParser = (<|> fail "invalid topic") $ Topic <$> do + void A.peekWord8' + level <- pLevel + levels <- A.many' (pSlash >> pLevel) + A.endOfInput + pure (level :| levels) + where + pSlash = void (A.word8 slash) + pLevel = Level . BS.toShort <$> A.takeWhile + (\w8-> w8 /= slash && w8 /= zero && w8 /= hash && w8 /= plus) +{-# INLINEABLE topicParser #-} + +topicBuilder :: Topic -> BS.Builder +topicBuilder (Topic (Level x:|xs)) = + foldl' + (\acc (Level l)-> acc <> BS.word8 slash <> BS.shortByteString l) + (BS.shortByteString x) xs +{-# INLINE topicBuilder #-} + +filterBuilder :: Filter -> BS.Builder +filterBuilder (Filter (Level x:|xs)) = + foldl' + (\acc (Level l)-> acc <> BS.word8 slash <> BS.shortByteString l) + (BS.shortByteString x) xs +{-# INLINE filterBuilder #-} + +topicLength :: Topic -> Int +topicLength (Topic (Level x:|xs)) = + BS.length x + len' xs 0 + where + len' [] acc = acc + len' (Level z:zs) acc = len' zs $! acc + 1 + BS.length z +{-# INLINE topicLength #-} + +filterLength :: Filter -> Int +filterLength (Filter (Level x:|xs)) = + BS.length x + len' xs 0 + where + len' [] acc = acc + len' (Level z:zs) acc = len' zs $! acc + 1 + BS.length z +{-# INLINE filterLength #-} + +filterParser :: A.Parser Filter +filterParser = (<|> fail "invalid filter") $ Filter <$> do + void A.peekWord8' + (x:xs) <- pLevels + pure (x:|xs) + where + pSlash = void (A.word8 slash) + pLevel = Level . BS.toShort <$> A.takeWhile + (\w8-> w8 /= slash && w8 /= zero && w8 /= hash && w8 /= plus) + pLevels + = (void (A.word8 hash) >> A.endOfInput >> pure [multiLevelWildcard]) + <|> (void (A.word8 plus) >> ((A.endOfInput >> pure [singleLevelWildcard]) <|> + (pSlash >> (:) <$> pure singleLevelWildcard <*> pLevels))) + <|> (pLevel >>= \x-> (x:) <$> ((A.endOfInput >> pure []) <|> (pSlash >> pLevels))) +{-# INLINEABLE filterParser #-} + +levelParser :: A.Parser Level +levelParser = do + x <- A.takeWhile (\w8-> w8 /= slash && w8 /= zero) + A.endOfInput + pure (Level $ BS.toShort x) +{-# INLINE levelParser #-} + +-- | The @#@ path element. It must only appear at the end of a `Filter`. +multiLevelWildcard :: Level +multiLevelWildcard = Level $ BS.pack $ pure hash +{-# INLINE multiLevelWildcard #-} + +-- | The @+@ path element. It may appear anywhere within a `Filter`. +singleLevelWildcard :: Level +singleLevelWildcard = Level $ BS.pack $ pure plus +{-# INLINE singleLevelWildcard #-} + +-- | Returns `True` iff the `Level` starts with @$@. +startsWithDollar :: Level -> Bool +startsWithDollar (Level bs) = + not (BS.null bs) && BS.index bs 0 == dollar +{-# INLINE startsWithDollar #-} + +zero, plus, hash, slash, dollar :: Word8 +zero = 0x00 +plus = 0x2b +hash = 0x23 +slash = 0x2f +dollar = 0x24 diff --git a/src/Network/MQTT/Trie.hs b/src/Network/MQTT/Trie.hs new file mode 100644 --- /dev/null +++ b/src/Network/MQTT/Trie.hs @@ -0,0 +1,406 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE BangPatterns #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.MQTT.Trie +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.MQTT.Trie ( + -- * Trie + Trie (..) + , TrieValue (..) + -- ** null + , null + -- ** empty + , empty + -- ** size + , size + -- ** sizeWith + , sizeWith + -- ** singleton + , singleton + -- ** matchTopic + , matchTopic + -- ** matchFilter + , matchFilter + -- ** lookup + , lookup + -- ** findMaxBounded + , findMaxBounded + -- ** insert + , insert + -- ** insertWith + , insertWith + -- ** insertFoldable + , insertFoldable + -- ** map + , map + -- ** mapMaybe + , mapMaybe + -- ** foldl' + , foldl' + -- ** delete + , delete + -- ** union + , union + -- ** unionWith + , unionWith + -- ** differenceWith + , differenceWith + ) where + +import Control.Applicative ((<|>)) +import qualified Data.Binary as B +import Data.Functor.Identity +import qualified Data.IntSet as IS +import qualified Data.List as L +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.Map.Strict as M +import Data.Maybe hiding (mapMaybe) +import Data.Monoid +import Prelude hiding (lookup, map, null) + +import Network.MQTT.Message.Topic + +-- | The `Trie` is a map-like data structure designed to hold elements +-- that can efficiently be queried according to the matching rules specified +-- by MQTT. +-- The primary purpose is to manage client subscriptions, but it can just +-- as well be used to manage permissions etc. +-- +-- The tree consists of nodes that may or may not contain values. The edges +-- are filter components. As some value types have the concept of a null +-- (i.e. an empty set) the `TrieValue` is a class defining the data +-- family `TrieNode`. This is a performance and size optimization to +-- avoid unnecessary boxing and case distinction. +newtype Trie a = Trie { branches :: M.Map Level (TrieNode a) } + +class TrieValue a where + data TrieNode a + node :: Trie a -> Maybe a -> TrieNode a + nodeNull :: a -> Bool + nodeTree :: TrieNode a -> Trie a + nodeValue :: TrieNode a -> Maybe a + +instance (TrieValue a, Monoid a) => Monoid (Trie a) where + mempty = empty + mappend = unionWith mappend + +instance (TrieValue a, Eq a) => Eq (Trie a) where + Trie m1 == Trie m2 = + M.size m1 == M.size m2 && and (zipWith f (M.toAscList m1) (M.toAscList m2)) + where + f (l1,n1) (l2,n2) = l1 == l2 && nodeValue n1 == nodeValue n2 && nodeTree n1 == nodeTree n2 + +instance (TrieValue a, Show a) => Show (Trie a) where + show (Trie m) = "Trie [" ++ L.intercalate ", " (f <$> M.toAscList m) ++ "]" + where + f (l,n) = "(" ++ show l ++ ", Node (" ++ show (nodeValue n) ++ ") (" ++ show (nodeTree n) ++ ")" + +instance B.Binary (Trie ()) where + put _ = pure () + get = pure empty + +empty :: Trie a +empty = Trie mempty + +null :: Trie a -> Bool +null (Trie m) = M.null m + +-- | Count all trie nodes that are not `nodeNull`. +size :: TrieValue a => Trie a -> Int +size = sizeWith (const 1) + +sizeWith :: TrieValue a => (a -> Int) -> Trie a -> Int +sizeWith sz = countTrie 0 + where + -- Depth-first search through the tree. + -- This implementation uses an accumulator in order to not defer + -- the evaluation of the additions. + countTrie !accum t = + M.foldl' countNode accum (branches t) + countNode !accum n = + case nodeValue n of + Nothing -> countTrie accum (nodeTree n) + Just v -> countTrie (accum + sz v) (nodeTree n) + +singleton :: TrieValue a => Filter -> a -> Trie a +singleton tf = singleton' (filterLevels tf) + where + singleton' (x:|xs) a + | nodeNull a = empty + | otherwise = Trie $ M.singleton x $ case xs of + [] -> node empty (Just a) + (y:ys) -> node (singleton' (y:|ys) a) Nothing + +insert :: TrieValue a => Filter -> a -> Trie a -> Trie a +insert = insertWith const + +insertWith :: TrieValue a => (a -> a -> a) -> Filter -> a -> Trie a -> Trie a +insertWith f tf a = insertWith' (filterLevels tf) + where + insertWith' (x:|xs) (Trie m) + | nodeNull a = Trie m + | otherwise = Trie $ M.alter g x m + where + g mn = Just $ case xs of + [] -> case mn of + Nothing -> node empty (Just a) + Just n -> node (nodeTree n) $ (f a <$> nodeValue n) <|> Just a + (y:ys) -> node (insertWith' (y:|ys) $ fromMaybe empty $ nodeTree <$> mn) Nothing + +insertFoldable :: (TrieValue a, Foldable t) => t (Filter, a) -> Trie a -> Trie a +insertFoldable = flip $ foldr $ uncurry insert + +delete :: TrieValue a => Filter -> Trie a -> Trie a +delete tf = delete' (filterLevels tf) + where + delete' (x:|xs) (Trie m) = Trie $ M.update g x m + where + g n = case xs of + [] | null (nodeTree n) -> Nothing + | otherwise -> Just $ node (nodeTree n) Nothing + y:ys -> let t = delete' (y:|ys) (nodeTree n) in + case nodeValue n of + Nothing | null t -> Nothing + | otherwise -> Just $ node t Nothing + Just v -> Just $ node t (Just v) + +map :: (TrieValue a, TrieValue b) => (a -> b) -> Trie a -> Trie b +map f (Trie m) = Trie $ fmap g m + where + g n = let t = map f (nodeTree n) in node t (f <$> nodeValue n) + +-- | Applies a functor to a try and removes nodes for which the mapping +-- function returns `Nothing`. +mapMaybe :: (TrieValue a, TrieValue b) => (a -> Maybe b) -> Trie a -> Trie b +mapMaybe f (Trie m) = Trie (M.mapMaybe g m) + where + g n | isNothing v' && null t' = Nothing + | otherwise = Just (node t' v') + where + v' = nodeValue n >>= f + t' = mapMaybe f $ nodeTree n + +foldl' :: (TrieValue b) => (a -> b -> a) -> a -> Trie b -> a +foldl' f acc (Trie m) = M.foldl' g acc m + where + g acc' n = flip (foldl' f) (nodeTree n) $! case nodeValue n of + Nothing -> acc' + Just value -> f acc' value + +union :: (TrieValue a, Monoid a) => Trie a -> Trie a -> Trie a +union (Trie m1) (Trie m2) = Trie (M.unionWith g m1 m2) + where + g n1 n2 = node (nodeTree n1 `union` nodeTree n2) (nodeValue n1 <> nodeValue n2) + +unionWith :: (TrieValue a) => (a -> a -> a) -> Trie a -> Trie a -> Trie a +unionWith f (Trie m1) (Trie m2) = Trie (M.unionWith g m1 m2) + where + g n1 n2 = node (unionWith f (nodeTree n1) (nodeTree n2)) (nodeValue n1 `merge` nodeValue n2) + merge (Just v1) (Just v2) = Just (f v1 v2) + merge mv1 mv2 = mv1 <|> mv2 + +differenceWith :: (TrieValue a, TrieValue b) => (a -> b -> Maybe a) -> Trie a -> Trie b -> Trie a +differenceWith f (Trie m1) (Trie m2) = Trie (M.differenceWith g m1 m2) + where + g n1 n2 = k (differenceWith f (nodeTree n1) (nodeTree n2)) (d (nodeValue n1) (nodeValue n2)) + d (Just v1) (Just v2) = f v1 v2 + d (Just v1) _ = Just v1 + d _ _ = Nothing + k t Nothing | null t = Nothing + | otherwise = Just $ node t Nothing + k t (Just v) | null t && nodeNull v = Nothing + | otherwise = Just $ node t $ Just v + +-- | Collect all values of nodes that match a given topic (according to the +-- matching rules specified by the MQTT protocol). +lookup :: (TrieValue a, Monoid a) => Topic -> Trie a -> a +lookup tf = fromMaybe mempty . lookupHead (topicLevels tf) + where + -- If the first level starts with $ then it must not be matched against + and #. + lookupHead (x:|xs) t@(Trie m) + | startsWithDollar x = case xs of + [] -> M.lookup x m >>= nodeValue + (y:ys) -> M.lookup x m >>= lookupTail y ys . nodeTree + | otherwise = lookupTail x xs t + lookupTail x [] (Trie m) = + matchSingleLevelWildcard <> matchMultiLevelWildcard <> matchComponent + where + matchSingleLevelWildcard = M.lookup singleLevelWildcard m >>= nodeValue + matchMultiLevelWildcard = M.lookup multiLevelWildcard m >>= nodeValue + matchComponent = M.lookup x m >>= \n-> + case M.lookup multiLevelWildcard $ branches $ nodeTree n of + -- component match, but no additional multiLevelWildcard below + Nothing -> nodeValue n + -- component match and multiLevelWildcard match below + Just n' -> nodeValue n <> nodeValue n' + lookupTail x (y:ys) (Trie m) = + matchSingleLevelWildcard <> matchMultiLevelWildcard <> matchComponent + where + matchSingleLevelWildcard = M.lookup singleLevelWildcard m >>= lookupTail y ys . nodeTree + matchMultiLevelWildcard = M.lookup multiLevelWildcard m >>= nodeValue + matchComponent = M.lookup x m >>= lookupTail y ys . nodeTree + +-- | Find the greatest value in a trie that matches the topic. +-- +-- * Stops search as soon as a `maxBound` element has been found. +-- * Doesn't match into `$` topics. +findMaxBounded :: (TrieValue a, Ord a, Bounded a) => Topic -> Trie a -> Maybe a +findMaxBounded topic = findHead (topicLevels topic) + where + findHead (x:|xs) t@(Trie m) + | startsWithDollar x = case xs of + [] -> M.lookup x m >>= nodeValue + (y:ys) -> M.lookup x m >>= findTail y ys . nodeTree + | otherwise = findTail x xs t + findTail x [] (Trie m) = + matchMultiLevelWildcard `maxBounded` matchSingleLevelWildcard `maxBounded` matchComponent + where + matchMultiLevelWildcard = M.lookup multiLevelWildcard m >>= + nodeValue + matchSingleLevelWildcard = M.lookup singleLevelWildcard m >>= \n-> + nodeValue n `maxBounded` (nodeValue =<< M.lookup multiLevelWildcard (branches $ nodeTree n)) + matchComponent = M.lookup x m >>= \n-> + nodeValue n `maxBounded` (nodeValue =<< M.lookup multiLevelWildcard (branches $ nodeTree n)) + findTail x (y:ys) (Trie m) = + matchMultiLevelWildcard `maxBounded` matchSingleLevelWildcard `maxBounded` matchComponent + where + matchMultiLevelWildcard = M.lookup multiLevelWildcard m >>= nodeValue + matchSingleLevelWildcard = M.lookup singleLevelWildcard m >>= findTail y ys . nodeTree + matchComponent = M.lookup x m >>= findTail y ys . nodeTree + + maxBounded :: (Ord a, Bounded a) => Maybe a -> Maybe a -> Maybe a + maxBounded a b + | a == Just maxBound = a + | otherwise = max a b + +-- | Match a `Topic` against a `Trie`. +-- +-- The function returns true iff the tree contains at least one node that +-- matches the topic /and/ contains a value (including nodes that are +-- indirectly matched by wildcard characters like `+` and `#` as described +-- in the MQTT specification). +matchTopic :: TrieValue a => Topic -> Trie a -> Bool +matchTopic tf = matchTopicHead (topicLevels tf) + where + -- The '#' is always a terminal node and therefore does not contain subtrees. + -- By invariant, a '#' node only exists if it contains a value. For this + -- reason it does not need to be checked for a value here, but just for + -- existence. + -- A '+' node on the other hand may contain subtrees and may not carry a value + -- itself. This needs to be checked. + matchTopicHead (x:|xs) t@(Trie m) + | startsWithDollar x = case xs of + [] -> matchExact x m + (y:ys) -> fromMaybe False $ matchTopicTail y ys . nodeTree <$> M.lookup x m + | otherwise = matchTopicTail x xs t + matchTopicTail x [] (Trie m) = + matchExact x m || matchPlus || matchHash + where + matchPlus = isJust ( nodeValue =<< M.lookup singleLevelWildcard m ) + matchHash = M.member multiLevelWildcard m + matchTopicTail x (y:ys) (Trie m) = + M.member multiLevelWildcard m || case M.lookup x m of + -- Same is true for '#' node here. In case no '#' hash node is present it is + -- first tried to match the exact topic and then to match any '+' node. + Nothing -> matchPlus + Just n -> matchTopicTail y ys (nodeTree n) || matchPlus + where + -- A '+' node matches any topic element. + matchPlus = fromMaybe False + $ matchTopicTail y ys . nodeTree <$> M.lookup singleLevelWildcard m + -- An exact match is the case if the map contains a node for the key and + -- the node is not empty _or_ the node's subtree contains a wildcard key (wildcards) + -- always also match the parent node. + matchExact x m = case M.lookup x m of + Nothing -> False + Just n -> isJust (nodeValue n) || let Trie m' = nodeTree n in M.member multiLevelWildcard m' + +-- | Match a `Filter` against a `Trie`. +-- +-- The function returns true iff the tree contains a path that is +-- /less or equally specific/ than the filter and the terminal node contains +-- a value that is not `nodeNull`. +-- +-- > match (singleton "#") "a" = True +-- > match (singleton "#") "+" = True +-- > match (singleton "#") "a/+/b" = True +-- > match (singleton "#") "a/+/#" = True +-- > match (singleton "+") "a" = True +-- > match (singleton "+") "+" = True +-- > match (singleton "+") "+/a" = False +-- > match (singleton "+") "#" = False +-- > match (singleton "a") "a" = True +-- > match (singleton "a") "b" = False +-- > match (singleton "a") "+" = False +-- > match (singleton "a") "#" = False +matchFilter :: TrieValue a => Filter -> Trie a -> Bool +matchFilter tf = matchFilter' (filterLevels tf) + where + matchFilter' (x:|[]) (Trie m) + | x == multiLevelWildcard = matchMultiLevelWildcard + | x == singleLevelWildcard = matchMultiLevelWildcard || matchSingleLevelWildcard + | otherwise = matchMultiLevelWildcard || matchSingleLevelWildcard || matchExact + where + matchMultiLevelWildcard = M.member multiLevelWildcard m + matchSingleLevelWildcard = isJust ( nodeValue =<< M.lookup singleLevelWildcard m ) + matchExact = case M.lookup x m of + Nothing -> False + Just n' -> isJust (nodeValue n') || let Trie m' = nodeTree n' in M.member multiLevelWildcard m' + matchFilter' (x:|y:zs) (Trie m) + | x == multiLevelWildcard = matchMultiLevelWildcard + | x == singleLevelWildcard = matchMultiLevelWildcard || matchSingleLevelWildcard + | otherwise = matchMultiLevelWildcard || matchSingleLevelWildcard || matchExact + where + matchMultiLevelWildcard = M.member multiLevelWildcard m + matchSingleLevelWildcard = fromMaybe False $ matchFilter' (y:|zs) . nodeTree <$> M.lookup singleLevelWildcard m + matchExact = fromMaybe False $ matchFilter' (y:|zs) . nodeTree <$> M.lookup x m + +-------------------------------------------------------------------------------- +-- Specialised nodeTree implemenations using data families +-------------------------------------------------------------------------------- + +instance TrieValue IS.IntSet where + data TrieNode IS.IntSet = IntSetTrieNode !(Trie IS.IntSet) !IS.IntSet + node t = IntSetTrieNode t . fromMaybe mempty + nodeNull = IS.null + nodeTree (IntSetTrieNode t _) = t + nodeValue (IntSetTrieNode _ v) + | nodeNull v = Nothing + | otherwise = Just v + +instance TrieValue (Identity a) where + data TrieNode (Identity a) = IdentityNode !(Trie (Identity a)) !(Maybe (Identity a)) + node t n@Nothing = IdentityNode t n + node t n@(Just _) = IdentityNode t n + nodeNull = const False + nodeTree (IdentityNode t _) = t + nodeValue (IdentityNode _ mv) = mv + +instance TrieValue () where + data TrieNode () = UnitNode {-# UNPACK #-} !Int !(Trie ()) + node t Nothing = UnitNode 0 t + node t _ = UnitNode 1 t + nodeNull = const False + nodeTree (UnitNode _ t) = t + nodeValue (UnitNode 0 _) = Nothing + nodeValue (UnitNode _ _) = Just () + +instance TrieValue Bool where + data TrieNode Bool = BoolNode {-# UNPACK #-} !Int !(Trie Bool) + node t Nothing = BoolNode 0 t + node t (Just False) = BoolNode 1 t + node t (Just True) = BoolNode 2 t + nodeNull = const False + nodeTree (BoolNode _ t) = t + nodeValue (BoolNode 1 _) = Just False + nodeValue (BoolNode 2 _) = Just True + nodeValue (BoolNode _ _) = Nothing diff --git a/src/Network/Stack/Server.hs b/src/Network/Stack/Server.hs new file mode 100644 --- /dev/null +++ b/src/Network/Stack/Server.hs @@ -0,0 +1,258 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +-------------------------------------------------------------------------------- +-- | +-- Module : Network.Stack.Server +-- Copyright : (c) Lars Petersen 2016 +-- License : MIT +-- +-- Maintainer : info@lars-petersen.net +-- Stability : experimental +-------------------------------------------------------------------------------- +module Network.Stack.Server where + +import Control.Concurrent.Async +import Control.Concurrent.MVar +import qualified Control.Exception as E +import Control.Monad +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BS +import qualified Data.ByteString.Builder.Extra as BS +import qualified Data.ByteString.Lazy as BSL +import Data.Int +import Data.Typeable +import qualified Data.X509 as X509 +import qualified Network.TLS as TLS +import qualified Network.WebSockets as WS +import qualified Network.WebSockets.Stream as WS +import qualified System.Socket as S +import qualified System.Socket.Type.Stream as S + +data TLS a +data WebSocket a + +class (Typeable a) => ServerStack a where + data Server a + data ServerConfig a + data ServerException a + data ServerConnection a + data ServerConnectionInfo a + -- | Creates a new server from a configuration and passes it to a handler function. + -- + -- The server given to the handler function shall be bound and in + -- listening state. The handler function is usually a + -- `Control.Monad.forever` loop that accepts and handles new connections. + -- + -- > withServer config $ \server-> + -- > forever $ withConnection handleConnection + withServer :: ServerConfig a -> (Server a -> IO b) -> IO b + -- | Waits for and accepts a new connection from a listening server and passes + -- it to a handler function. + -- + -- This operation is blocking until the lowest layer in the stack accepts + -- a new connection. The handlers of all other layers are executed within + -- an `Control.Concurrent.Async.Async` which is returned. This allows + -- the main thread waiting on the underlying socket to block just as long + -- as necessary. Upper layer protocol handshakes (TLS etc) will be executed + -- in the new thread. + -- + -- > withServer config $ \server-> forever $ + -- > future <- withConnection server handleConnection + -- > putStrLn "The lowest layer accepted a new connection!" + -- > async $ do + -- > result <- wait future + -- > putStrLn "The connection handler returned:" + -- > print result + withConnection :: Server a -> (ServerConnection a -> ServerConnectionInfo a -> IO b) -> IO (Async b) + +class ServerStack a => StreamServerStack a where + sendStream :: ServerConnection a -> BS.ByteString -> IO Int + sendStream server bs = fromIntegral <$> sendStreamLazy server (BSL.fromStrict bs) + sendStreamLazy :: ServerConnection a -> BSL.ByteString -> IO Int64 + sendStreamLazy server = foldM + (\sent bs-> sendStream server bs >>= \sent'-> pure $! sent + fromIntegral sent') 0 . BSL.toChunks + sendStreamBuilder :: ServerConnection a -> Int -> BS.Builder -> IO Int64 + sendStreamBuilder server chunksize = sendStreamLazy server + . BS.toLazyByteStringWith (BS.untrimmedStrategy chunksize chunksize) mempty + receiveStream :: ServerConnection a -> Int -> IO BS.ByteString + receiveStream server i = BSL.toStrict <$> receiveStreamLazy server i + receiveStreamLazy :: ServerConnection a -> Int -> IO BSL.ByteString + receiveStreamLazy server i = BSL.fromStrict <$> receiveStream server i + {-# MINIMAL (sendStream|sendStreamLazy), (receiveStream|receiveStreamLazy) #-} + +-- | This class is an abstraction for `ServerStack`s that support the +-- transmission and reception of finite messages. +class ServerStack a => MessageServerStack a where + type ClientMessage a + type ServerMessage a + -- | Send a message. + -- + -- - Returns the encoded message size. + sendMessage :: ServerConnection a -> ServerMessage a -> IO Int64 + -- | Send several messages. This might lead to an improvement for very short messages. + -- + -- - Returns the summed size of all encoded messages. + sendMessages :: Foldable t => ServerConnection a -> t (ServerMessage a) -> IO Int64 + -- | Receive a message. + -- + -- - The second parameter determines the maximum encoded message size which + -- must not be exceeded by the client or an exception will be thrown. + -- Implementations shall track the consumed bytes and shall throw an + -- exception as soon as the limit is exceeded even if the message is not + -- yet complete. This is important to prevent _denial of service_ attacks. + receiveMessage :: ServerConnection a -> Int64 -> IO (ClientMessage a) + -- | Consumes incoming messages with a supplied consumer callback. + -- + -- - The second parameter limits the size of a single encoded message + -- (see `receiveMessage`). + consumeMessages :: ServerConnection a -> Int64 -> (ClientMessage a -> IO Bool) -> IO () + +instance (Typeable f, Typeable p, S.Family f, S.Protocol p) => StreamServerStack (S.Socket f S.Stream p) where + sendStream (SocketServerConnection s) bs = S.sendAll s bs S.msgNoSignal + sendStreamLazy (SocketServerConnection s) lbs = S.sendAllLazy s lbs S.msgNoSignal + sendStreamBuilder (SocketServerConnection s) bufsize builder = S.sendAllBuilder s bufsize builder S.msgNoSignal + receiveStream (SocketServerConnection s) i = S.receive s i S.msgNoSignal + +instance (StreamServerStack a) => StreamServerStack (TLS a) where + sendStreamLazy connection lbs = TLS.sendData (tlsContext connection) lbs >> pure (BSL.length lbs) + receiveStream connection _ = TLS.recvData (tlsContext connection) + +instance (StreamServerStack a) => StreamServerStack (WebSocket a) where + sendStream connection bs = WS.sendBinaryData (wsConnection connection) bs >> pure (BS.length bs) + sendStreamLazy connection lbs = WS.sendBinaryData (wsConnection connection) lbs >> pure (BSL.length lbs) + receiveStreamLazy connection _ = WS.receiveData (wsConnection connection) + +instance (S.Family f, S.Type t, S.Protocol p, Typeable f, Typeable t, Typeable p) => ServerStack (S.Socket f t p) where + data Server (S.Socket f t p) = SocketServer + { socketServer :: !(S.Socket f t p) + , socketServerConfig :: !(ServerConfig (S.Socket f t p)) + } + data ServerConfig (S.Socket f t p) = SocketServerConfig + { socketServerConfigBindAddress :: !(S.SocketAddress f) + , socketServerConfigListenQueueSize :: Int + } + data ServerException (S.Socket f t p) = SocketServerException !S.SocketException + data ServerConnection (S.Socket f t p) = SocketServerConnection !(S.Socket f t p) + data ServerConnectionInfo (S.Socket f t p) = SocketServerConnectionInfo !(S.SocketAddress f) + withServer c handle = E.bracket + (SocketServer <$> S.socket <*> pure c) + (S.close . socketServer) $ \server-> do + S.setSocketOption (socketServer server) (S.ReuseAddress True) + S.bind (socketServer server) (socketServerConfigBindAddress $ socketServerConfig server) + S.listen (socketServer server) (socketServerConfigListenQueueSize $ socketServerConfig server) + handle server + withConnection server handle = + E.bracketOnError (S.accept (socketServer server)) (S.close . fst) $ \(connection, addr)-> + async (handle (SocketServerConnection connection) (SocketServerConnectionInfo addr) `E.finally` S.close connection) + +instance (StreamServerStack a, Typeable a) => ServerStack (TLS a) where + data Server (TLS a) = TlsServer + { tlsTransportServer :: Server a + , tlsServerConfig :: ServerConfig (TLS a) + } + data ServerConfig (TLS a) = TlsServerConfig + { tlsTransportConfig :: ServerConfig a + , tlsServerParams :: TLS.ServerParams + } + data ServerException (TLS a) = + TlsServerEndOfStreamException + deriving (Eq, Ord, Show) + data ServerConnection (TLS a) = TlsServerConnection + { tlsTransportConnection :: ServerConnection a + , tlsContext :: TLS.Context + } + data ServerConnectionInfo (TLS a) = TlsServerConnectionInfo + { tlsTransportServerConnectionInfo :: ServerConnectionInfo a + , tlsCertificateChain :: Maybe X509.CertificateChain + } + withServer config handle = + withServer (tlsTransportConfig config) $ \server-> + handle (TlsServer server config) + withConnection server handle = + withConnection (tlsTransportServer server) $ \connection info-> do + let backend = TLS.Backend { + TLS.backendFlush = pure () -- backend doesn't buffer + , TLS.backendClose = pure () -- backend gets closed automatically + , TLS.backendSend = void . sendStream connection + -- The following is problematic: The TLS implementation requires us + -- to return exactly as many bytes as requested. The underlying transport + -- though only yields as many bytes as available. + -- The solution is to read, append and loop until the request + -- can be fulfilled. + -- TODO: Use bytestring builder for concatenation. + -- TODO: Fix TLS library upstream. The interface is awkward for a + -- networking lib. + , TLS.backendRecv = flip (receiveExactly connection) mempty + } + mvar <- newEmptyMVar + let srvParams = tlsServerParams $ tlsServerConfig server + srvParams' = srvParams { + TLS.serverHooks = (TLS.serverHooks srvParams) { + TLS.onClientCertificate = \certChain-> do + putMVar mvar certChain + pure TLS.CertificateUsageAccept + } + } + context <- TLS.contextNew backend srvParams' + TLS.handshake context + certificateChain <- tryTakeMVar mvar + x <- handle + (TlsServerConnection connection context) + (TlsServerConnectionInfo info certificateChain) + TLS.bye context + pure x + where + receiveExactly connection bytes accum = do + bs <- receiveStream connection bytes + -- TCP sockets signal a graceful end of the stream by returning zero bytes. + -- This function is not allowed to return less than the bytes + -- request and we shall not loop forever here (we did!). There is no + -- other option than throwing an exception here. + when (BS.null bs) $ + E.throwIO (TlsServerEndOfStreamException :: ServerException (TLS a)) + if BS.length bs < bytes + then receiveExactly connection (bytes - BS.length bs) $! accum `mappend` bs + else pure $! accum `mappend` bs + +instance (StreamServerStack a) => ServerStack (WebSocket a) where + data Server (WebSocket a) = WebSocketServer + { wsTransportServer :: Server a + } + data ServerConfig (WebSocket a) = WebSocketServerConfig + { wsTransportConfig :: ServerConfig a + } + data ServerException (WebSocket a) = WebSocketServerException + data ServerConnection (WebSocket a) = WebSocketServerConnection + { wsTransportConnection :: ServerConnection a + , wsConnection :: WS.Connection + } + data ServerConnectionInfo (WebSocket a) = WebSocketServerConnectionInfo + { wsTransportServerConnectionInfo :: ServerConnectionInfo a + , wsRequestHead :: WS.RequestHead + } + withServer config handle = + withServer (wsTransportConfig config) $ \server-> + handle (WebSocketServer server) + withConnection server handle = + withConnection (wsTransportServer server) $ \connection info-> do + let readSocket = (\bs-> if BS.null bs then Nothing else Just bs) <$> receiveStream connection 4096 + let writeSocket Nothing = pure () + writeSocket (Just bs) = void (sendStream connection (BSL.toStrict bs)) + stream <- WS.makeStream readSocket writeSocket + pendingConnection <- WS.makePendingConnectionFromStream stream (WS.ConnectionOptions $ pure ()) + acceptedConnection <- WS.acceptRequestWith pendingConnection (WS.AcceptRequest $ Just "mqtt") + x <- handle + (WebSocketServerConnection connection acceptedConnection) + (WebSocketServerConnectionInfo info $ WS.pendingRequest pendingConnection) + WS.sendClose acceptedConnection ("Thank you for flying Haskell." :: BS.ByteString) + pure x + +deriving instance Show (S.SocketAddress f) => Show (ServerConnectionInfo (S.Socket f t p)) +deriving instance Show (ServerConnectionInfo a) => Show (ServerConnectionInfo (TLS a)) +deriving instance Show (ServerConnectionInfo a) => Show (ServerConnectionInfo (WebSocket a)) + +instance Typeable a => E.Exception (ServerException (TLS a)) diff --git a/test/BrokerTest.hs b/test/BrokerTest.hs new file mode 100644 --- /dev/null +++ b/test/BrokerTest.hs @@ -0,0 +1,327 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +module BrokerTest ( getTestTree ) where + +import Control.Concurrent +import Control.Concurrent.Async +import Control.Exception +import Data.Monoid +import Data.String +import Control.Monad +import qualified Data.Sequence as Seq +import Data.Typeable +import Data.UUID (UUID) +import Test.Tasty +import Test.Tasty.HUnit + +import qualified Network.MQTT.Broker as Broker +import Network.MQTT.Broker.Authentication +import qualified Network.MQTT.Broker.Session as Session +import Network.MQTT.Message +import qualified Network.MQTT.Message as Message +import qualified Network.MQTT.Trie as R + +newtype TestAuthenticator = TestAuthenticator (AuthenticatorConfig TestAuthenticator) + +instance Authenticator TestAuthenticator where + data AuthenticatorConfig TestAuthenticator = TestAuthenticatorConfig + { cfgAuthenticate :: ConnectionRequest -> IO (Maybe UUID) + , cfgGetPrincipal :: UUID -> IO (Maybe Principal) + } + data AuthenticationException TestAuthenticator = TestAuthenticatorException deriving (Typeable, Show) + newAuthenticator = pure . TestAuthenticator + authenticate (TestAuthenticator cfg) req = cfgAuthenticate cfg req + getPrincipal (TestAuthenticator cfg) uuid = cfgGetPrincipal cfg uuid + +instance Exception (AuthenticationException TestAuthenticator) + +getTestTree :: IO TestTree +getTestTree = + pure $ testGroup "Broker" + [ testGroup "Authentication" + [ testCase "Reject with 'ServerUnavaible' when authentication throws exception" $ do + m1 <- newEmptyMVar + m2 <- newEmptyMVar + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigNoService + let sessionRejectHandler = putMVar m1 + sessionAcceptHandler session present = putMVar m2 (session, present) + Broker.withSession broker connectionRequest sessionRejectHandler sessionAcceptHandler + tryReadMVar m1 >>= \x-> Just ServerUnavailable @?= x + tryReadMVar m2 >>= \x-> Nothing @?= x + , testCase "Reject 'NotAuthorized' when authentication returned Nothing" $ do + m1 <- newEmptyMVar + m2 <- newEmptyMVar + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigNoAccess + let sessionRejectHandler = putMVar m1 + sessionAcceptHandler session present = putMVar m2 (session, present) + Broker.withSession broker connectionRequest sessionRejectHandler sessionAcceptHandler + tryReadMVar m1 >>= \x-> Just NotAuthorized @?= x + tryReadMVar m2 >>= \x-> Nothing @?= x + ] + , testGroup "Subscriptions" + + [ testCase "subscribe the same filter from 2 different sessions" $ do + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + let req1 = connectionRequest { requestClientIdentifier = "1" } + req2 = connectionRequest { requestClientIdentifier = "2" } + msg = Message.Message "a/b" QoS0 (Retain False) "" + Broker.withSession broker req1 (const $ pure ()) $ \session1 _ -> + Broker.withSession broker req2 (const $ pure ()) $ \session2 _ -> do + Session.subscribe session1 (PacketIdentifier 42) [("a/b", QoS0)] + Session.subscribe session2 (PacketIdentifier 47) [("a/b", QoS0)] + Broker.publishDownstream broker msg + queue1 <- (<>) <$> Session.dequeue session1 <*> Session.dequeue session1 + queue2 <- (<>) <$> Session.dequeue session2 <*> Session.dequeue session2 + queue1 @?= Seq.fromList [ ServerSubscribeAcknowledged (PacketIdentifier 42) [Just QoS0], ServerPublish (PacketIdentifier (-1)) (Duplicate False) msg] + queue2 @?= Seq.fromList [ ServerSubscribeAcknowledged (PacketIdentifier 47) [Just QoS0], ServerPublish (PacketIdentifier (-1)) (Duplicate False) msg] + + , testCase "get retained message on subscription (newer overrides older, issue #6)" $ do + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + let req1 = connectionRequest { requestClientIdentifier = "1" } + msg1 = Message.Message "topic" QoS0 (Retain True) "test" + msg2 = Message.Message "topic" QoS0 (Retain True) "toast" + Broker.withSession broker req1 (const $ pure ()) $ \session1 _-> do + Broker.publishDownstream broker msg1 + Broker.publishDownstream broker msg2 + Session.subscribe session1 (PacketIdentifier 23) [("topic", QoS0)] + queue1 <- (<>) <$> Session.dequeue session1 <*> Session.dequeue session1 + queue1 @?= Seq.fromList [ ServerSubscribeAcknowledged (PacketIdentifier 23) [Just QoS0], ServerPublish (PacketIdentifier (-1)) (Duplicate False) msg2] + + , testCase "delete retained message when body is empty" $ do + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + let req1 = connectionRequest { requestClientIdentifier = "1" } + msg1 = Message.Message "topic" QoS0 (Retain True) "test" + msg2 = Message.Message "topic" QoS0 (Retain True) "" + Broker.withSession broker req1 (const $ pure ()) $ \session1 _ -> do + Broker.publishDownstream broker msg1 + Broker.publishDownstream broker msg2 + Session.subscribe session1 (PacketIdentifier 23) [("topic", QoS0)] + queue1 <- (<>) <$> Session.dequeue session1 <*> Session.dequeue session1 + queue1 @?= Seq.fromList [ ServerSubscribeAcknowledged (PacketIdentifier 23) [Just QoS0] ] + ] + + , testGroup "Queue overflow handling" + + [ testCase "Barrel shift on overflowing QoS0 queue" $ do + let msgs = [ Message.Message "topic" QoS0 (Retain False) (fromString $ show x) | x <- [(1 :: Int)..] ] + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + t1 <- newEmptyMVar + t2 <- newEmptyMVar + t3 <- newEmptyMVar + t4 <- newEmptyMVar + t5 <- newEmptyMVar + let h session _ = do { + Session.subscribe session (PacketIdentifier 0) [("topic", QoS0)]; + putMVar t1 (); + takeMVar t2; + void $ Session.dequeue session; -- subscribe acknowledge + putMVar t3 =<< Session.dequeue session; + takeMVar t4; + putMVar t5 =<< Session.dequeue session; + } + let w = Broker.withSession broker connectionRequest (const $ pure ()) h + withAsync w $ \_-> do + takeMVar t1 + forM_ (take 10 msgs) $ Broker.publishDownstream broker + putMVar t2 () + p <- takeMVar t3 + assertEqual "Expect 10 packets being dequeued (first time)." + (Seq.fromList $ fmap (ServerPublish (PacketIdentifier (-1)) (Duplicate False)) $ take 10 msgs) p + forM_ (take 11 msgs) $ Broker.publishDownstream broker + putMVar t4 () + q <- takeMVar t5 + assertEqual "Expect 10 packets being dequeued (second time)." + (Seq.fromList $ fmap (ServerPublish (PacketIdentifier (-1)) (Duplicate False)) $ take 10 $ drop 1 msgs) q + + , testCase "Terminate session on overflowing QoS1 queue" $ do + let msgs = [ Message.Message "topic" QoS1 (Retain False) (fromString $ show x) | x <- [(1 :: Int)..] ] + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + t1 <- newEmptyMVar + t2 <- newEmptyMVar + t3 <- newEmptyMVar + t4 <- newEmptyMVar + t5 <- newEmptyMVar + let h session _ = do { + Session.subscribe session (PacketIdentifier 0) [("topic", QoS1)]; + putMVar t1 (); + takeMVar t2; + putMVar t3 =<< (Seq.drop 1 <$> Session.dequeue session); -- cut off subscribe acknowledge + void $ takeMVar t4; + putMVar t5 =<< Session.dequeue session; + } + let w = Broker.withSession broker connectionRequest (const $ pure ()) h + withAsync w $ \as-> do + takeMVar t1 + forM_ (take 10 msgs) $ Broker.publishDownstream broker + putMVar t2 () + p <- takeMVar t3 + assertEqual "Expect 10 packets being dequeued (first time)." + (Seq.fromList $ zipWith (\i m-> ServerPublish (PacketIdentifier i) (Duplicate False) m) [0..] $ take 10 msgs) p + forM_ (take 11 msgs) $ Broker.publishDownstream broker + assertEqual "Expect session handler thread to be killed." "Left thread killed" =<< (show <$> waitCatch as) + + , testCase "overflowing Qos2 queue (session termination)" $ do + let msgs = [ Message.Message "topic" QoS2 (Retain False) (fromString $ show x) | x <- [(1 :: Int)..] ] + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + t1 <- newEmptyMVar + t2 <- newEmptyMVar + t3 <- newEmptyMVar + t4 <- newEmptyMVar + t5 <- newEmptyMVar + let h session _ = do { + Session.subscribe session (PacketIdentifier 0) [("topic", QoS2)]; + putMVar t1 (); + takeMVar t2; + putMVar t3 =<< (Seq.drop 1 <$> Session.dequeue session); -- cut off subscribe acknowledge + void $ takeMVar t4; + putMVar t5 =<< Session.dequeue session; + } + let w = Broker.withSession broker connectionRequest (const $ pure ()) h + withAsync w $ \as-> do + takeMVar t1 + forM_ (take 10 msgs) $ Broker.publishDownstream broker + putMVar t2 () + p <- takeMVar t3 + assertEqual "Expect 10 packets being dequeued (first time)." + (Seq.fromList $ zipWith (\i m-> ServerPublish (PacketIdentifier i) (Duplicate False) m) [0..] $ take 10 msgs) p + forM_ (take 11 msgs) $ Broker.publishDownstream broker + assertEqual "Expect session handler thread to be killed." "Left thread killed" =<< (show <$> waitCatch as) + ] + + , testGroup "Quality of Service" + + [ testCase "transmit a QoS1 message and process acknowledgement" $ do + let msg = Message.Message "topic" QoS1 (Retain False) "payload" + pid = PacketIdentifier 0 + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + Broker.withSession broker connectionRequest (const $ pure ()) $ \session _-> do + pids1 <- Session.getFreePacketIdentifiers session + Session.enqueueMessage session msg + queue <- Session.dequeue session + pids2 <- Session.getFreePacketIdentifiers session + Session.processPublishAcknowledged session pid + pids3 <- Session.getFreePacketIdentifiers session + assertEqual "One packet identifier shall be in use after `dequeue`." (Seq.drop 1 pids1) pids2 + assertEqual "The packet identifier shall have been returned after the message has been acknowledged." pids1 pids3 + assertEqual "The packet is expected in the output queue." (Seq.fromList [ ServerPublish pid (Duplicate False) msg ]) queue + + , testCase "receive a QoS1 message and send acknowledgement" $ do + let msg = Message.Message "topic" QoS1 (Retain False) "payload" + pid = PacketIdentifier 0 + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + Broker.withSession broker connectionRequest (const $ pure ()) $ \session _-> do + Session.subscribe session pid [("topic", QoS0)] + queue1 <- Session.dequeue session + assertEqual "A subscribe acknowledgement shall be in the output queue." (Seq.fromList [ ServerSubscribeAcknowledged pid [Just QoS0] ]) queue1 + Session.processPublish session pid (Duplicate False) msg + queue2 <- Session.dequeue session + assertEqual "A publish acknowledgment and the (downgraded) message itself shall be in the output queue." (Seq.fromList [ ServerPublishAcknowledged pid ]) queue2 + queue3 <- Session.dequeue session + assertEqual "The downgraded message queue shall be in the output queue." (Seq.fromList [ ServerPublish (PacketIdentifier (-1)) (Duplicate False) msg { msgQoS = QoS0} ]) queue3 + + , testCase "transmit a QoS1 message and retransmit after connection failure" $ do + let req = connectionRequest { requestCleanSession = False } + msg = Message.Message "topic" QoS1 (Retain False) "payload" + pid = PacketIdentifier 0 + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + Broker.withSession broker req (const $ pure ()) $ \session present-> do + assertEqual "The session shall not be present." (SessionPresent False) present + Session.enqueueMessage session msg + queue <- Session.dequeue session + assertEqual "The message shall be in the output queue." (Seq.fromList [ ServerPublish pid (Duplicate False) msg]) queue + Broker.withSession broker req (const $ pure ()) $ \session present-> do + assertEqual "The session shall be present." (SessionPresent True) present + queue <- Session.dequeue session + assertEqual "The message shall again be in the output queue, and must not be marked duplicate." (Seq.fromList [ ServerPublish pid (Duplicate True) msg ]) queue + + , testCase "transmit a QoS2 message and process confirmations" $ do + let msg = Message.Message "topic" QoS2 (Retain False) "payload" + pid = PacketIdentifier 0 + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + Broker.withSession broker connectionRequest (const $ pure ()) $ \session _-> do + pids1 <- Session.getFreePacketIdentifiers session + Session.enqueueMessage session msg + queue2 <- Session.dequeue session + pids2 <- Session.getFreePacketIdentifiers session + assertEqual "One packet identifier shall be in use after `dequeue`." (Seq.drop 1 pids1) pids2 + assertEqual "A PUB packet is expected in the output queue." (Seq.fromList [ ServerPublish pid (Duplicate False) msg ]) queue2 + Session.processPublishReceived session pid + queue3 <- Session.dequeue session + pids3 <- Session.getFreePacketIdentifiers session + assertEqual "The packet identifier shall still be in use." (Seq.drop 1 pids1) pids3 + assertEqual "A PUBREL packet is expected in the output queue." (Seq.fromList [ ServerPublishRelease pid ]) queue3 + Session.processPublishComplete session pid + pids4 <- Session.getFreePacketIdentifiers session + assertEqual "The packet identifier shall have been returned after the transaction has been completed." pids1 pids4 + + , testCase "transmit a QoS2 message and handle retransmissions on connection failure" $ do + let req = connectionRequest { requestCleanSession = False } + msg = Message.Message "topic" QoS2 (Retain False) "payload" + pid = PacketIdentifier 0 + broker <- Broker.newBroker $ TestAuthenticator authenticatorConfigAllAccess + Broker.withSession broker req (const $ pure ()) $ \session _-> do + Session.enqueueMessage session msg + queue <- Session.dequeue session + assertEqual "The message shall be in the output queue." (Seq.fromList [ ServerPublish pid (Duplicate False) msg ]) queue + Broker.withSession broker req (const $ pure ()) $ \session _-> do + queue <- Session.dequeue session + assertEqual "The message shall again be in the output queue and must be marked duplicate." (Seq.fromList [ ServerPublish pid (Duplicate True) msg ]) queue + Session.processPublishReceived session pid + queue' <- Session.dequeue session + assertEqual "The release command shall be in the output queue." (Seq.fromList [ ServerPublishRelease pid ]) queue' + Broker.withSession broker req (const $ pure ()) $ \session _-> do + queue <- Session.dequeue session + assertEqual "The release command shall be in the output queue (again)." (Seq.fromList [ ServerPublishRelease pid ]) queue + Session.processPublishComplete session pid + Broker.withSession broker req (const $ pure ()) $ \session _-> do + queue <- Session.dequeue session + assertEqual "The output queue shall be empty." mempty queue + ] + ] + +authenticatorConfigNoService :: AuthenticatorConfig TestAuthenticator +authenticatorConfigNoService = TestAuthenticatorConfig + { cfgAuthenticate = const $ throwIO TestAuthenticatorException + , cfgGetPrincipal = const $ throwIO TestAuthenticatorException + } + +authenticatorConfigNoAccess :: AuthenticatorConfig TestAuthenticator +authenticatorConfigNoAccess = TestAuthenticatorConfig + { cfgAuthenticate = const (pure Nothing) + , cfgGetPrincipal = const (pure Nothing) + } + +authenticatorConfigAllAccess :: AuthenticatorConfig TestAuthenticator +authenticatorConfigAllAccess = TestAuthenticatorConfig + { cfgAuthenticate = const $ pure (Just uuid) + , cfgGetPrincipal = const $ pure (Just pcpl) + } + where + uuid = read "3c7efc50-bff0-4e09-9a9b-0f2bff2db8fc" + pcpl = Principal { + principalUsername = Nothing + , principalQuota = quota + , principalPublishPermissions = R.singleton "#" () + , principalSubscribePermissions = R.singleton "#" () + , principalRetainPermissions = R.singleton "#" () + } + quota = Quota { + quotaMaxIdleSessionTTL = 60 + , quotaMaxPacketSize = 65535 + , quotaMaxPacketIdentifiers = 10 + , quotaMaxQueueSizeQoS0 = 10 + , quotaMaxQueueSizeQoS1 = 10 + , quotaMaxQueueSizeQoS2 = 10 + } + +connectionRequest :: ConnectionRequest +connectionRequest = ConnectionRequest + { requestClientIdentifier = "mqtt-default" + , requestCleanSession = True + , requestSecure = False + , requestCredentials = Nothing + , requestHttp = Nothing + , requestCertificateChain = Nothing + , requestRemoteAddress = Nothing + } diff --git a/test/EncodingTest.hs b/test/EncodingTest.hs new file mode 100644 --- /dev/null +++ b/test/EncodingTest.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE OverloadedStrings #-} +module EncodingTest (tests) where + +import qualified Data.Binary.Get as SG +import qualified Data.ByteString.Builder as BS +import qualified Data.ByteString.Lazy as BSL +import Prelude hiding (head) +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck as QC + +import Network.MQTT.Message +import qualified Network.MQTT.Message as Topic + +tests :: TestTree +tests = testGroup "Encoding / Decoding" + [ testClientPacket + , testServerPacket + , testLengthEncoding + ] + +testClientPacket :: TestTree +testClientPacket = QC.testProperty "clientMessageBuilder <-> clientPacketParser" $ \msg-> + msg === SG.runGet clientPacketParser (BS.toLazyByteString $ clientPacketBuilder msg) + +testServerPacket :: TestTree +testServerPacket = QC.testProperty "serverMessageBuilder <-> serverPacketParser" $ \msg-> + msg === SG.runGet serverPacketParser (BS.toLazyByteString $ serverPacketBuilder msg) + +instance Arbitrary ClientPacket where + arbitrary = oneof + [ arbitraryConnect + , arbitraryPublish + , arbitrarySubscribe + , arbitraryUnsubscribe + , ClientPublishAcknowledged <$> arbitrary + , ClientPublishReceived <$> arbitrary + , ClientPublishRelease <$> arbitrary + , ClientPublishComplete <$> arbitrary + , pure ClientPingRequest + , pure ClientDisconnect + ] + where + arbitraryConnect = ClientConnect + <$> elements [ "", "client-identifier" ] + <*> arbitrary + <*> arbitrary + <*> arbitrary + <*> oneof [ pure Nothing, (Just .) . (,) + <$> elements [ "", "username" ] + <*> oneof [ pure Nothing, Just <$> elements [ Password "", Password "password" ] ] ] + arbitraryPublish = do + msg <- arbitrary + dup <- Duplicate <$> if msgQoS msg == QoS0 then pure False else arbitrary + pid <- PacketIdentifier <$> if msgQoS msg == QoS0 then pure (-1) else choose (0, 65535) + pure (ClientPublish pid dup msg) + arbitrarySubscribe = ClientSubscribe + <$> arbitrary + <*> listOf1 ((,) <$> arbitrary <*> arbitrary ) + arbitraryUnsubscribe = ClientUnsubscribe + <$> arbitrary + <*> listOf1 arbitrary + +instance Arbitrary ServerPacket where + arbitrary = oneof + [ ServerConnectionAccepted <$> arbitrary + , ServerConnectionRejected <$> arbitrary + , do + msg <- arbitrary + dup <- Duplicate <$> if msgQoS msg == QoS0 then pure False else arbitrary + pid <- PacketIdentifier <$> if msgQoS msg == QoS0 then pure (-1) else choose (0, 65535) + pure (ServerPublish pid dup msg) + , ServerPublishAcknowledged <$> arbitrary + , ServerPublishReceived <$> arbitrary + , ServerPublishRelease <$> arbitrary + , ServerPublishComplete <$> arbitrary + , ServerSubscribeAcknowledged <$> arbitrary <*> listOf1 arbitrary + , ServerUnsubscribeAcknowledged <$> arbitrary + , pure ServerPingResponse + ] + +instance Arbitrary Message where + arbitrary = Message + <$> arbitrary + <*> arbitrary + <*> arbitrary + <*> elements [ "", "shortTopic", Payload $ BSL.replicate 345 23 ] + +instance Arbitrary PacketIdentifier where + arbitrary = PacketIdentifier <$> choose (0, 65535) + +instance Arbitrary SessionPresent where + arbitrary = SessionPresent <$> arbitrary + +instance Arbitrary CleanSession where + arbitrary = CleanSession <$> arbitrary + +instance Arbitrary Retain where + arbitrary = Retain <$> arbitrary + +instance Arbitrary KeepAliveInterval where + arbitrary = KeepAliveInterval <$> arbitrary + +instance Arbitrary RejectReason where + arbitrary = elements + [ UnacceptableProtocolVersion + , IdentifierRejected + , ServerUnavailable + , BadUsernameOrPassword + , NotAuthorized + ] + +instance Arbitrary QoS where + arbitrary = elements [ QoS0, QoS1, QoS2 ] + +instance Arbitrary Topic.Topic where + arbitrary = elements + [ "a" + , "/" + , "//" + , "a/b/c" + , "a/nyːnɔʃk/c" + , " /" + , " / /" + , "$SYS/grampfglompf" + ] + +instance Arbitrary Topic.Filter where + arbitrary = elements + [ "#" + , "+/#" + , "a/a/a/a/a/a/a/a/adasdas//asda/+/+/ /#" + , "$SYS/#" + ] + +testLengthEncoding :: TestTree +testLengthEncoding = + testGroup "lengthParser, lengthBuilder" + [ testCase "p [193,2] == 321" $ assertEqual "" 321 + ( SG.runGet lengthParser $ BSL.pack [193,2] ) + + , testCase "p [0x00] == 0" $ assertEqual "" 0 + ( SG.runGet lengthParser (BSL.pack [0x00]) ) + + , testCase "p [0x7f] == 127" $ assertEqual "" 127 + ( SG.runGet lengthParser (BSL.pack [0x7f]) ) + + , testCase "p [0x80, 0x01] == 128" $ assertEqual "" 128 + ( SG.runGet lengthParser (BSL.pack [0x80, 0x01]) ) + + , testCase "p [0xff, 0x7f] == 16383" $ assertEqual "" 16383 + ( SG.runGet lengthParser (BSL.pack [0xff, 0x7f]) ) + + , testCase "p [0x80, 0x80, 0x01] == 16384" $ assertEqual "" 16384 + ( SG.runGet lengthParser (BSL.pack [0x80, 0x80, 0x01]) ) + + , testCase "p [0xff, 0xff, 0x7f] == 2097151" $ assertEqual "" 2097151 + ( SG.runGet lengthParser (BSL.pack [0xff, 0xff, 0x7f]) ) + + , testCase "p [0x80, 0x80, 0x80, 0x01] == 2097152" $ assertEqual "" 2097152 + ( SG.runGet lengthParser (BSL.pack [0x80, 0x80, 0x80, 0x01]) ) + + , testCase "p [0xff, 0xff, 0xff, 0x7f] == 268435455" $ assertEqual "" 268435455 + ( SG.runGet lengthParser (BSL.pack [0xff, 0xff, 0xff, 0x7f]) ) + + , testCase "p [0xff, 0xff, 0xff, 0xff] == invalid" $ assertEqual "" + ( Left ("", 4, "lengthParser: invalid input") ) + ( SG.runGetOrFail lengthParser (BSL.pack [0xff, 0xff, 0xff, 0xff]) ) + + , QC.testProperty "lengthParser . lengthBuilder == id" $ + \i -> let i' = i `mod` 268435455 + in i' == SG.runGet lengthParser (BS.toLazyByteString (lengthBuilder i')) + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,18 @@ +module Main where + +import Test.Tasty + +import qualified BrokerTest +import qualified EncodingTest +import qualified TrieTest +import qualified TopicTest + +main :: IO () +main = do + brokerTests <- BrokerTest.getTestTree + defaultMain $ testGroup "Network.MQTT" + [ TrieTest.tests + , TopicTest.tests + , EncodingTest.tests + , brokerTests + ] diff --git a/test/PrioritySemaphoreTest.hs b/test/PrioritySemaphoreTest.hs new file mode 100644 --- /dev/null +++ b/test/PrioritySemaphoreTest.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE OverloadedStrings, LambdaCase #-} +module Main where + +import Control.Monad ( void ) +import Control.Exception +import Control.Concurrent +import Control.Concurrent.Async +import Test.Tasty +import Test.Tasty.HUnit + +import qualified Control.Concurrent.PrioritySemaphore as PrioritySemaphore + +main :: IO () +main = defaultMain $ testGroup "Control.Concurrent.PrioritySemaphore" + [ testCase "2 threads, subsequently" $ do + sem <- PrioritySemaphore.newPrioritySemaphore + m1 <- newEmptyMVar :: IO (MVar Int) + m2 <- newEmptyMVar :: IO (MVar Int) + PrioritySemaphore.exclusively sem $ + putMVar m1 1 + PrioritySemaphore.exclusively sem $ + putMVar m2 2 + assertEqual "m1" 1 =<< takeMVar m1 + assertEqual "m2" 2 =<< takeMVar m2 + + , testCase "2 threads, one interrupting the other" $ do + sem <- PrioritySemaphore.newPrioritySemaphore + m1 <- newEmptyMVar :: IO (MVar Int) + m2 <- newEmptyMVar :: IO (MVar Int) + m3 <- newEmptyMVar :: IO (MVar Int) + t1 <- async $ + PrioritySemaphore.exclusively sem $ do + putMVar m1 1 + threadDelay 2000000 `finally` putMVar m2 3 + void $ swapMVar m1 2 + pure (4 :: Int) + t2 <- async $ do + threadDelay 1000000 + PrioritySemaphore.exclusively sem $ do + putMVar m3 5 + pure (6 :: Int) + threadDelay 3000000 + waitCatch t1 >>= \case + Left _ -> pure () + Right _ -> assertFailure "thread 1 should have been interrupted" + waitCatch t2 >>= \case + Left _ -> assertFailure "thread 2 should not have failed" + Right x -> assertEqual "t2" 6 x + assertEqual "m1" 1 =<< takeMVar m1 + assertEqual "m2" 3 =<< takeMVar m2 + assertEqual "m3" 5 =<< takeMVar m3 + ] diff --git a/test/RetainedStoreStrictnessTest.hs b/test/RetainedStoreStrictnessTest.hs new file mode 100644 --- /dev/null +++ b/test/RetainedStoreStrictnessTest.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main where + +import Control.Monad (forM_) + +import qualified Network.MQTT.Message as M +import qualified Network.MQTT.Broker.RetainedMessages as Retained + +-- | This shall test whether inserting into the RetainedStore +-- leaks memory by building up unevaluated thunks. +-- +-- The test is supposed to be run with '+RTS -M4m'. +-- In case of a memory leak the program will crash as +-- the thunks would require around 300MB heap. +main :: IO () +main = do + store <- Retained.new + forM_ [1..1000000 :: Int] $ \_i-> do + Retained.store message store + where + message :: M.Message + message = M.Message { + M.msgTopic = "ahsdjkha/def/hij" + , M.msgPayload = "ahsjdkhajskdhaksjdhakjshd" + , M.msgQoS = M.QoS1 + , M.msgRetain = M.Retain True + } diff --git a/test/TopicTest.hs b/test/TopicTest.hs new file mode 100644 --- /dev/null +++ b/test/TopicTest.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE OverloadedStrings #-} +module TopicTest ( tests ) where + +import qualified Data.Attoparsec.ByteString as A +import Data.List.NonEmpty (NonEmpty (..)) +import Test.Tasty +import Test.Tasty.HUnit + +import Network.MQTT.Message + +tests :: TestTree +tests = testGroup "Topic / Filter / Level" + [ testGroup "Topic" + [ testGroup "show / fromString" + [ testCase "\"/\"" $ assertEqual "" "\"/\"" $ show ("/" :: Topic) + , testCase "\"\x2603\"" $ assertEqual "" "\"\x2603\"" $ show ("\x2603" :: Topic) + , testCase "\"a/b/c\"" $ assertEqual "" "\"a/b/c\"" $ show ("a/b/c" :: Topic) + ] + , testGroup "topicParser" + [ testCase "! \"\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "" + , testCase "! \"\\NUL\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "\NUL" + , testCase "! \"+\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "+" + , testCase "! \"#\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "#" + , testCase " \"/\"" $ assertEqual "" (Right $ "":|[""]) $ topicLevels <$> A.parseOnly topicParser "/" + , testCase " \"//\"" $ assertEqual "" (Right $ "":|["",""]) $ topicLevels <$> A.parseOnly topicParser "//" + , testCase " \"/a\"" $ assertEqual "" (Right $ "":|["a"]) $ topicLevels <$> A.parseOnly topicParser "/a" + , testCase " \"a\"" $ assertEqual "" (Right $ "a":|[]) $ topicLevels <$> A.parseOnly topicParser "a" + , testCase " \"a/\"" $ assertEqual "" (Right $ "a":|[""]) $ topicLevels <$> A.parseOnly topicParser "a/" + , testCase " \"a/bcd\"" $ assertEqual "" (Right $ "a":|["bcd"]) $ topicLevels <$> A.parseOnly topicParser "a/bcd" + , testCase "! \"a/b+d\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "a/b+d" + , testCase "! \"a/b#d\"" $ assertEqual "" (Left "Failed reading: invalid topic") $ topicLevels <$> A.parseOnly topicParser "a/b#d" + ] + ] + , testGroup "Filter" + [ testGroup "show / fromString" + [ testCase "\"/\"" $ assertEqual "" "\"/\"" $ show ("/" :: Filter) + , testCase "\"\x2603\"" $ assertEqual "" "\"\x2603\"" $ show ("\x2603" :: Filter) + , testCase "\"a/b/c\"" $ assertEqual "" "\"a/b/c\"" $ show ("a/b/c" :: Filter) + , testCase "\"#\"" $ assertEqual "" "\"#\"" $ show ("#" :: Filter) + , testCase "\"+\"" $ assertEqual "" "\"+\"" $ show ("+" :: Filter) + , testCase "\"a/#\"" $ assertEqual "" "\"a/#\"" $ show ("a/#" :: Filter) + , testCase "\"+/+/#\"" $ assertEqual "" "\"+/+/#\"" $ show ("+/+/#" :: Filter) + , testCase "\"/#\"" $ assertEqual "" "\"/#\"" $ show ("/#" :: Filter) + ] + , testGroup "filterParser" + [ testCase "! \"\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "" + , testCase "! \"\\NUL\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "\NUL" + , testCase " \"+\"" $ assertEqual "" (Right $ "+":|[]) $ filterLevels <$> A.parseOnly filterParser "+" + , testCase " \"#\"" $ assertEqual "" (Right $ "#":|[]) $ filterLevels <$> A.parseOnly filterParser "#" + , testCase "! \"#/\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "#/" + , testCase "! \"a/a+b\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "a/a+b" + , testCase "! \"a/a#b\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "a/a#b" + , testCase "! \"#a\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "#a" + , testCase " \"/\"" $ assertEqual "" (Right $ "":|[""]) $ filterLevels <$> A.parseOnly filterParser "/" + , testCase " \"//\"" $ assertEqual "" (Right $ "":|["",""]) $ filterLevels <$> A.parseOnly filterParser "//" + , testCase " \"/a\"" $ assertEqual "" (Right $ "":|["a"]) $ filterLevels <$> A.parseOnly filterParser "/a" + , testCase " \"a\"" $ assertEqual "" (Right $ "a":|[]) $ filterLevels <$> A.parseOnly filterParser "a" + , testCase " \"a/\"" $ assertEqual "" (Right $ "a":|[""]) $ filterLevels <$> A.parseOnly filterParser "a/" + , testCase " \"a/b\"" $ assertEqual "" (Right $ "a":|["b"]) $ filterLevels <$> A.parseOnly filterParser "a/b" + , testCase " \"a/+/c123/#\"" $ assertEqual "" (Right $ "a":|["+","c123","#"]) $ filterLevels <$> A.parseOnly filterParser "a/+/c123/#" + , testCase "! \"a/+/c123/#/d\"" $ assertEqual "" (Left "Failed reading: invalid filter") $ filterLevels <$> A.parseOnly filterParser "a/+/c123/#/d" + ] + ] + ] diff --git a/test/TrieSizeTest.hs b/test/TrieSizeTest.hs new file mode 100644 --- /dev/null +++ b/test/TrieSizeTest.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main where + +import Control.Monad ( foldM, when ) +import Data.IntSet as IS +import qualified Data.List.NonEmpty as NL +import qualified Data.Map as M +import System.Random ( randomIO ) + +import Network.MQTT.Message +import Network.MQTT.Trie as R + +-- | This test shall assure that the memory consumption of the subscription +-- tree stays within certain limits. + +-- This test suite runs with restricted maximum heap size. +-- The routing tree with 1 million nodes and around 3 subscriptions per node +-- took 83MB when measured. +-- We allow for some deviation, but it should not exceed 100MB heap space. +main :: IO () +main = do + r <- randomTree 6 10 :: IO (Trie IS.IntSet) + -- We need to to something with `r` or it won't be evaluated. + when (R.size r == 0) (error "should not be 0") + +randomTree :: Int -> Int -> IO (Trie IS.IntSet) +randomTree 0 _ = Trie <$> pure mempty +randomTree depth branching = Trie <$> foldM (\m e-> + flip (M.insert e) m <$> (R.node + <$> randomTree (depth - 1) branching + <*> randomSet :: IO (TrieNode IS.IntSet))) M.empty + (take branching randomTreeElements) + where + randomSet :: IO (Maybe IS.IntSet) + randomSet = Just <$> f 0 IS.empty + where + f :: Int -> IS.IntSet -> IO IS.IntSet + f i accum = do + p <- randomIO :: IO Double + if p < 0.25 + then pure accum + else f (succ i) $! IS.insert i accum + +randomTreeElements :: [Level] +randomTreeElements = + fmap (NL.head . topicLevels) ["a","b","c","d","e","f","g","h","i","j","k","l","m" :: Topic] diff --git a/test/TrieTest.hs b/test/TrieTest.hs new file mode 100644 --- /dev/null +++ b/test/TrieTest.hs @@ -0,0 +1,227 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +module TrieTest ( tests ) where + +import Data.Functor.Identity +import qualified Data.IntSet as IS +import Prelude hiding (head) + +import Test.Tasty +import Test.Tasty.HUnit + +import qualified Network.MQTT.Trie as R + +tests :: TestTree +tests = testGroup "Trie" + [ testGroup "null" + [ testCase "! null $ singleton \"a\" ()" $ assertBool "" $ not $ R.null $ R.singleton "a" () + , testCase "! null $ singleton \"a/b\" ()" $ assertBool "" $ not $ R.null $ R.singleton "a/b" () + ] + , testGroup "empty" + [ testCase " null empty" $ assertBool "not null" $ R.null R.empty + ] + , testGroup "size" + [ testCase "size empty == 0" $ R.size (R.empty :: R.Trie ()) @?= 0 + , testCase "size tree1 == 9" $ R.size tree1 @?= 9 + , testCase "size tree2 == 2" $ R.size tree2 @?= 2 + , testCase "size tree3 == 2" $ R.size tree3 @?= 2 + , testCase "size tree4 == 3" $ R.size tree4 @?= 3 + ] + , testGroup "sizeWith" + [ testCase "sizeWith IS.size empty == 0" $ R.sizeWith IS.size R.empty @?= 0 + , testCase "sizeWith IS.size tree1 == 9" $ R.sizeWith IS.size tree1 @?= 9 + , testCase "sizeWith IS.size tree2 == 2" $ R.sizeWith IS.size tree2 @?= 2 + , testCase "sizeWith IS.size tree3 == 7" $ R.sizeWith IS.size tree3 @?= 7 + ] + , testGroup "singleton" + [ testCase " matchTopic \"a\" $ singleton \"a\" ()" $ assertBool "" $ R.matchTopic "a" $ R.singleton "a" () + , testCase " matchTopic \"a/b/c\" $ singleton \"a/b/c\" ()" $ assertBool "" $ R.matchTopic "a/b/c" $ R.singleton "a/b/c" () + ] + , testGroup "matchTopic" + [ testCase " matchTopic \"a\" $ singleton \"a\" ()" $ assertBool "" $ R.matchTopic "a" $ R.singleton "a" () + , testCase " matchTopic \"a\" $ singleton \"#\" ()" $ assertBool "" $ R.matchTopic "a" $ R.singleton "#" () + , testCase " matchTopic \"a\" $ singleton \"a/#\" ()" $ assertBool "" $ R.matchTopic "a" $ R.singleton "a/#" () + , testCase " matchTopic \"a/b\" $ singleton \"a/#\" ()" $ assertBool "" $ R.matchTopic "a/b" $ R.singleton "a/#" () + , testCase " matchTopic \"a/b/c\" $ singleton \"a/#\" ()" $ assertBool "" $ R.matchTopic "a/b/c" $ R.singleton "a/#" () + , testCase "! matchTopic \"b/c/d\" $ singleton \"a/#\" ()" $ assertBool "" $ not $ R.matchTopic "b/c/d" $ R.singleton "a/#" () + , testCase "! matchTopic \"a\" $ singleton \"a/+\" ()" $ assertBool "" $ not $ R.matchTopic "a" $ R.singleton "a/+" () + , testCase "! matchTopic \"a\" $ singleton \"/a\" ()" $ assertBool "" $ not $ R.matchTopic "a" $ R.singleton "/a" () + , testCase " matchTopic \"a/b\" $ singleton \"a/b\" ()" $ assertBool "" $ R.matchTopic "a/b" $ R.singleton "a/b" () + , testCase " matchTopic \"a/b\" $ singleton \"a/+\" ()" $ assertBool "" $ R.matchTopic "a/b" $ R.singleton "a/+" () + , testCase " matchTopic \"a/b\" $ singleton \"a/#\" ()" $ assertBool "" $ R.matchTopic "a/b" $ R.singleton "a/#" () + , testCase " matchTopic \"a/b\" $ singleton \"a/b/#\" ()" $ assertBool "" $ R.matchTopic "a/b" $ R.singleton "a/b/#" () + , testCase "! matchTopic \"$SYS\" $ singleton \"#\" ()" $ assertBool "" $ not $ R.matchTopic "$SYS" $ R.singleton "#" () + , testCase "! matchTopic \"$SYS/a\" $ singleton \"#\" ()" $ assertBool "" $ not $ R.matchTopic "$SYS/a" $ R.singleton "#" () + , testCase " matchTopic \"$SYS\" $ singleton \"$SYS/#\" ()" $ assertBool "" $ R.matchTopic "$SYS" $ R.singleton "$SYS/#" () + , testCase " matchTopic \"$SYS/a\" $ singleton \"$SYS/#\" ()" $ assertBool "" $ R.matchTopic "$SYS/a" $ R.singleton "$SYS/#" () + , testCase "! matchTopic \"$SYS\" $ singleton \"+\" ()" $ assertBool "" $ not $ R.matchTopic "$SYS" $ R.singleton "+" () + , testCase "! matchTopic \"$SYS/a\" $ singleton \"+\" ()" $ assertBool "" $ not $ R.matchTopic "$SYS/a" $ R.singleton "+" () + , testCase "! matchTopic \"$SYS\" $ singleton \"$SYS/+\" ()" $ assertBool "" $ not $ R.matchTopic "$SYS" $ R.singleton "$SYS/+" () + , testCase " matchTopic \"$SYS/a\" $ singleton \"$SYS/+\" ()" $ assertBool "" $ R.matchTopic "$SYS/a" $ R.singleton "$SYS/+" () + ] + , testGroup "matchFilter" + [ testCase " matchFiler \"#\" $ singleton \"#\" ()" $ assertBool "" $ R.matchFilter "#" $ R.singleton "#" () + , testCase " matchFiler \"+\" $ singleton \"#\" ()" $ assertBool "" $ R.matchFilter "+" $ R.singleton "#" () + , testCase " matchFiler \"a\" $ singleton \"#\" ()" $ assertBool "" $ R.matchFilter "a" $ R.singleton "#" () + , testCase "! matchFiler \"#\" $ singleton \"+\" ()" $ assertBool "" $ not $ R.matchFilter "#" $ R.singleton "+" () + , testCase " matchFiler \"+\" $ singleton \"+\" ()" $ assertBool "" $ R.matchFilter "+" $ R.singleton "+" () + , testCase " matchFiler \"a\" $ singleton \"+\" ()" $ assertBool "" $ R.matchFilter "a" $ R.singleton "+" () + , testCase "! matchFiler \"#\" $ singleton \"a\" ()" $ assertBool "" $ not $ R.matchFilter "#" $ R.singleton "a" () + , testCase "! matchFiler \"+\" $ singleton \"a\" ()" $ assertBool "" $ not $ R.matchFilter "+" $ R.singleton "a" () + , testCase " matchFiler \"a\" $ singleton \"a\" ()" $ assertBool "" $ R.matchFilter "a" $ R.singleton "a" () + , testCase " matchFiler \"a\" $ singleton \"a/#\" ()" $ assertBool "" $ R.matchFilter "a" $ R.singleton "a/#" () + , testCase "! matchFiler \"a\" $ singleton \"a/+\" ()" $ assertBool "" $ not $ R.matchFilter "a" $ R.singleton "a/+" () + , testCase " matchFiler \"a/#\" $ singleton \"#\" ()" $ assertBool "" $ R.matchFilter "a/#" $ R.singleton "#" () + , testCase " matchFiler \"a/b/#\" $ singleton \"#\" ()" $ assertBool "" $ R.matchFilter "a/b/#" $ R.singleton "#" () + , testCase "! matchFiler \"a/+/c\" $ singleton \"a/b/c\" ()" $ assertBool "" $ not $ R.matchFilter "a/+/c" $ R.singleton "a/b/c" () + ] + , testGroup "lookup" + [ testCase "lookup \"a\" tree1 == [0,1,2,4]" $ R.lookup "a" tree1 @?= IS.fromList [0,1,2,4] + , testCase "lookup \"b\" tree1 == [0,1]" $ R.lookup "b" tree1 @?= IS.fromList [0,1] + , testCase "lookup \"a/a\" tree1 == [0,3,4]" $ R.lookup "a/a" tree1 @?= IS.fromList [0,3,4] + , testCase "lookup \"a/a/a\" tree1 == [0,4]" $ R.lookup "a/a/a" tree1 @?= IS.fromList [0,4] + , testCase "lookup \"$SYS\" tree1 == []" $ R.lookup "$SYS" tree1 @?= IS.fromList [] + , testCase "lookup \"$SYS/a\" tree1 == [5,6,7]" $ R.lookup "$SYS/a" tree1 @?= IS.fromList [5,6,7] + , testCase "lookup \"$SYS/a/a\" tree1 == [5,7,8]" $ R.lookup "$SYS/a/a" tree1 @?= IS.fromList [5,7,8] + ] + , testGroup "findMaxBounded" + [ testCase "findMaxBounded \"a/b/c\" tree5 == Just LT" $ R.findMaxBounded "a/b/c" tree5 @?= Just (Identity LT) + , testCase "findMaxBounded \"a/b/c/d\" tree5 == Just EQ" $ R.findMaxBounded "a/b/c/d" tree5 @?= Just (Identity EQ) + , testCase "findMaxBounded \"s\" tree5 == Just GT" $ R.findMaxBounded "s" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"t/a\" tree5 == Just GT" $ R.findMaxBounded "t/a" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"u/a\" tree5 == Just GT" $ R.findMaxBounded "u/a" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"x/a\" tree5 == Just GT" $ R.findMaxBounded "x/a" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"y/a\" tree5 == Just GT" $ R.findMaxBounded "y/a" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"z/x/x\" tree5 == Nothing" $ R.findMaxBounded "z/x/x" tree5 @?= Nothing + , testCase "findMaxBounded \"z/x/x/q\" tree5 == Just EQ" $ R.findMaxBounded "z/x/x/q" tree5 @?= Just (Identity EQ) + , testCase "findMaxBounded \"z/x/x/q/t\" tree5 == Nothing" $ R.findMaxBounded "z/x/x/q/t" tree5 @?= Nothing + , testCase "findMaxBounded \"z/x/r/q\" tree5 == Just GT" $ R.findMaxBounded "z/x/r/q" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"z/x/r/q/t\" tree5 == Just GT" $ R.findMaxBounded "z/x/r/q/t" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"SYS/SYS\" tree5 == Just GT" $ R.findMaxBounded "SYS/SYS" tree5 @?= Just (Identity GT) + , testCase "findMaxBounded \"$SYS/SYS\" tree5 == Nothing" $ R.findMaxBounded "$SYS/SYS" tree5 @?= Nothing + ] + , testGroup "insert" + [ testCase "lookup \"a/b\" tree2 == [3]" $ R.lookup "a/b" tree2 @?= IS.fromList [3] + , testCase "lookup \"a/b/c\" tree2 == [2]" $ R.lookup "a/b/c" tree2 @?= IS.fromList [2] + ] + , testGroup "insertWith" [ + testCase "parameter order (new, old)" $ ( R.lookup "a" $ R.insertWith const "a" (IS.singleton 2) $ R.singleton "a" (IS.singleton 1) ) @?= IS.fromList [2] + ] + , testGroup "map" [ ] + , testGroup "mapMaybe" [ + testCase "Trie (Identity [Int]) -> Trie (): #1" $ + let trie = R.insert "abc/#" (Identity [1 :: Int]) $ R.singleton "abc/def" (Identity [2]) + f (Identity xs) + | 1 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.singleton "abc/#" () + , testCase "Trie (Identity [Int]) -> Trie (): #2" $ + let trie = R.insert "abc/#" (Identity [1 :: Int]) $ R.singleton "abc/def" (Identity [2]) + f (Identity xs) + | 2 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.singleton "abc/def" () + , testCase "Trie (Identity [Int]) -> Trie (): #3" $ + let trie = R.insert "abc/#" (Identity [1 :: Int]) $ R.singleton "abc/def" (Identity [2]) + f (Identity xs) + | 3 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.empty + , testCase "Trie (Identity [Int]) -> Trie (): #4" $ + let trie = R.insert "abc" (Identity [1 :: Int]) $ R.insert "abc/hij" (Identity [2]) $ R.singleton "abc/def" (Identity [3]) + f (Identity xs) + | 1 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.singleton "abc" () + , testCase "Trie (Identity [Int]) -> Trie (): #5" $ + let trie = R.insert "abc" (Identity [1 :: Int]) $ R.insert "abc/hij" (Identity [2]) $ R.singleton "abc/def" (Identity [3]) + f (Identity xs) + | 2 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.singleton "abc/hij" () + , testCase "Trie (Identity [Int]) -> Trie (): #6" $ + let trie = R.insert "abc" (Identity [1 :: Int]) $ R.insert "abc/def" (Identity [2]) $ R.singleton "abc/def/hij" (Identity [1]) + f (Identity xs) + | 1 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= (R.insert "abc" () $ R.singleton "abc/def/hij" ()) + , testCase "Trie (Identity [Int]) -> Trie (): #7" $ + let trie = R.insert "abc" (Identity [1 :: Int]) $ R.insert "abc/def" (Identity [2]) $ R.singleton "abc/def/hij" (Identity [1]) + f (Identity xs) + | 2 `elem` xs = Just () + | otherwise = Nothing + in R.mapMaybe f trie @?= R.singleton "abc/def" () + ] + , testGroup "adjust" [ ] + , testGroup "delete" [ ] + , testGroup "union" [ + testCase "structurally distinct trees with shared prefix" + $ (R.singleton "a/b/y" $ IS.singleton 1) `R.union` (R.singleton "a/b/x" $ IS.singleton 2) + @?= R.insertFoldable [("a/b/y", IS.singleton 1), ("a/b/x", IS.singleton 2)] R.empty + , testCase "structurally equal trees with distinct values" + $ (R.singleton "a/b/x" $ IS.singleton 1) `R.union` (R.singleton "a/b/x" $ IS.singleton 2) + @?= R.singleton "a/b/x" (IS.fromList [1,2]) + ] + , testGroup "unionWith" [ + testCase "structurally distinct trees with shared prefix" + $ R.unionWith IS.union (R.singleton "a/b/y" $ IS.singleton 1) (R.singleton "a/b/x" $ IS.singleton 2) + @?= R.insertFoldable [("a/b/y", IS.singleton 1), ("a/b/x", IS.singleton 2)] R.empty + , testCase "structurally equal trees with distinct values" + $ R.unionWith IS.union (R.singleton "a/b/x" $ IS.singleton 1) (R.singleton "a/b/x" $ IS.singleton 2) + @?= R.singleton "a/b/x" (IS.fromList [1,2]) + ] + , testGroup "differenceWith" [ ] + ] + +tree1 :: R.Trie IS.IntSet +tree1 + = R.insertWith IS.union "#" (IS.singleton 0) + $ R.insertWith IS.union "+" (IS.singleton 1) + $ R.insertWith IS.union "a" (IS.singleton 2) + $ R.insertWith IS.union "a/+" (IS.singleton 3) + $ R.insertWith IS.union "a/#" (IS.singleton 4) + $ R.insertWith IS.union "$SYS/#" (IS.singleton 5) + $ R.insertWith IS.union "$SYS/+" (IS.singleton 6) + $ R.insertWith IS.union "$SYS/a/#" (IS.singleton 7) + $ R.insertWith IS.union "$SYS/+/a" (IS.singleton 8) + $ R.empty + +tree2 :: R.Trie IS.IntSet +tree2 + = R.insert "a/b" (IS.singleton 3) + $ R.insert "a/b/c" (IS.singleton 2) + $ R.insert "a/b/c" (IS.singleton 1) + $ R.empty + +tree3 :: R.Trie IS.IntSet +tree3 + = R.insert "a/b" (IS.fromList []) + $ R.insert "a/b/c" (IS.fromList [1,2,3,4]) + $ R.insert "a/b/d" (IS.fromList [5,6,7]) + $ R.insert "a/b/d/e" (IS.fromList []) + $ R.empty + +tree4 :: R.Trie () +tree4 + = R.insert "a/b" () + $ R.insert "a/b/c/d" () + $ R.insert "b/c/d" () + $ R.empty + +tree5 :: R.Trie (Identity Ordering) +tree5 + = R.insert "a/b/c" (Identity LT) + $ R.insert "a/b/c/d" (Identity EQ) + $ R.insert "s" (Identity EQ) + $ R.insert "s/#" (Identity GT) + $ R.insert "t/a" (Identity EQ) + $ R.insert "t/a/#" (Identity GT) + $ R.insert "u/+" (Identity EQ) + $ R.insert "u/a/#" (Identity GT) + $ R.insert "x/+" (Identity GT) + $ R.insert "x/a" (Identity undefined) + $ R.insert "y/#" (Identity GT) + $ R.insert "y/a" (Identity undefined) + $ R.insert "z/+/+/q" (Identity EQ) + $ R.insert "z/+/r/+/#" (Identity GT) + $ R.singleton "+/SYS" (Identity GT)