net-mqtt-0.8.8.0: src/Network/MQTT/Client.hs
{-|
Module : Network.MQTT.Client.
Description : An MQTT client.
Copyright : (c) Dustin Sallings, 2019
License : BSD3
Maintainer : dustin@spy.net
Stability : experimental
An MQTT protocol client
Both
<http://docs.oasis-open.org/mqtt/mqtt/v3.1.1/os/mqtt-v3.1.1-os.html MQTT 3.1.1>
and
<https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html MQTT 5.0>
are supported over plain TCP, TLS, WebSockets and Secure WebSockets.
-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Network.MQTT.Client (
-- * Configuring the client.
MQTTConfig(..), MQTTClient, QoS(..), Topic, mqttConfig, mkLWT, LastWill(..),
ProtocolLevel(..), Property(..), SubOptions(..), subOptions, MessageCallback(..),
-- * Running and waiting for the client.
waitForClient,
connectURI, isConnected,
disconnect, normalDisconnect, stopClient,
-- * General client interactions.
subscribe, unsubscribe, publish, publishq, pubAliased,
svrProps, connACK, MQTTException(..),
-- * Low-level bits
runMQTTConduit, MQTTConduit, isConnectedSTM, connACKSTM,
registerCorrelated, unregisterCorrelated
) where
import Control.Concurrent (myThreadId, threadDelay)
import Control.Concurrent.Async (Async, AsyncCancelled, async, asyncThreadId, cancel, cancelWith, link, race_, wait, waitCatch,
waitAnyCancel)
import Control.Concurrent.STM (STM, TChan, TQueue, TVar, atomically, check, modifyTVar', newTChan,
newTChanIO, newTQueueIO, newTVarIO, orElse, readTChan, readTQueue, readTVar,
readTVarIO, registerDelay, retry, writeTChan, writeTQueue, writeTVar)
import Control.DeepSeq (force)
import qualified Control.Exception as E
import Control.Monad (forever, guard, join, unless, void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Bifunctor (first)
import qualified Data.ByteString.Char8 as BCS
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Char8 as BC
import Data.Conduit (ConduitT, Void, await, runConduit, yield, (.|))
import Data.Conduit.Attoparsec (conduitParser)
import qualified Data.Conduit.Combinators as C
import Data.Conduit.Network (AppData, appSink, appSource, clientSettings, runTCPClient)
import Data.Conduit.Network.TLS (runTLSClient, tlsClientConfig, tlsClientTLSSettings)
import Data.Default.Class (def)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromJust, fromMaybe, isJust, isNothing)
import Data.Text (Text)
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Lazy.Encoding as TEL
import Data.Word (Word16)
import GHC.Conc (labelThread)
import Network.Connection (ConnectionParams (..), TLSSettings (..), connectTo, connectionClose,
connectionGetChunk, connectionPut, initConnectionContext)
import Network.URI (URI (..), nullURIAuth, unEscapeString, uriPort, uriRegName, uriUserInfo)
import qualified Network.WebSockets as WS
import Network.WebSockets.Stream (makeStream)
import System.IO.Error (catchIOError, isEOFError)
import System.Timeout (timeout)
import qualified Data.Set as Set
import Network.MQTT.Topic (Filter, Topic, mkTopic, unFilter, unTopic)
import Network.MQTT.Types as T
data ConnState = Starting
| Connected
| Stopping -- ^ Shutting down on request (DISCONNECT queued, or stopClient)
| Stopped
| Disconnected
| DiscoErr DisconnectRequest
| ConnErr ConnACKFlags deriving (Eq, Show)
-- | Where an inbound QoS 2 message is. It's delivered on PUBLISH, and
-- the PUBCOMP waits for both the PUBREL and the callback to finish.
data InboundQoS2 = Delivering -- ^ callback running, no PUBREL yet
| Released -- ^ PUBREL received, callback still running
| Delivered -- ^ callback done, waiting for PUBREL
deriving (Eq, Show)
data DispatchType = DSubACK | DUnsubACK | DPubACK | DPubREC | DPubREL | DPubCOMP
deriving (Eq, Show, Ord, Enum, Bounded)
-- | Callback invoked on each incoming subscribed message.
data MessageCallback = NoCallback
-- | Callbacks will be invoked asynchronously, ordering is likely preserved, but not guaranteed.
-- In high throughput scenarios, slow callbacks may result in a high number of Haskell threads,
-- potentially bringing down the entire application when running out of memory.
-- Typically faster than `OrderedCallback`.
-- Use 'System.Mem.setAllocationCounter' and 'System.Mem.enableAllocationLimit'
-- to restrain the unbounded threads memory use.
-- Use a simple 'System.Timeout.timeout' to restrain unbounded thread runtime.
| SimpleCallback (MQTTClient -> Topic -> BL.ByteString -> [Property] -> IO ())
-- | Callbacks are guaranteed to be invoked in the same order messages are received.
-- Slow callbacks do not block the underlying TCP connection. Pending QoS 0 messages
-- can accumulate while a callback is running. With MQTT 5, pending QoS 1 and 2
-- messages are bounded by the Receive Maximum. Messages received before the
-- connection ends are still delivered after it ends, unless 'stopClient' is called.
-- Typically slower than `SimpleCallback`.
| OrderedCallback (MQTTClient -> Topic -> BL.ByteString -> [Property] -> IO ())
-- | A LowLevelCallback receives the client and the entire publish request, providing
-- access to all of the fields of the request. This is slightly harder to use than
-- SimpleCallback for common cases, but there are cases where you need all the things.
| LowLevelCallback (MQTTClient -> PublishRequest -> IO ())
-- | A low level callback that is ordered.
| OrderedLowLevelCallback (MQTTClient -> PublishRequest -> IO ())
-- | The MQTT client.
--
-- See 'connectURI' for the most straightforward example.
data MQTTClient = MQTTClient {
_ch :: TChan MQTTPkt
, _pktID :: TVar Word16
, _cb :: MessageCallback
, _acks :: TVar (Map (DispatchType,Word16) (TChan MQTTPkt))
, _inflight :: TVar (Map Word16 InboundQoS2) -- ^ Inbound QoS 2 packet IDs awaiting PUBREL
, _st :: TVar ConnState
, _ct :: TVar (Maybe (Async ()))
, _outA :: TVar (Map Topic Word16)
, _inA :: TVar (Map Word16 Topic)
, _connACKFlags :: TVar ConnACKFlags
, _corr :: TVar (Map BL.ByteString MessageCallback)
, _cbQ :: TQueue (IO ())
, _cbHandle :: TVar (Maybe (Async ()))
, _outFlight :: TVar Word16 -- ^ QoS 1 and 2 PUBLISHes in flight to the server
}
-- | Configuration for setting up an MQTT client.
data MQTTConfig = MQTTConfig{
-- | False if a session should be reused.
-- This makes MQTT queue messages if the device goes offline.
-- Note you've to set '_connID' for this to be meaningful.
-- Brokers may assign one for you giving the wrong behavior.
-- AWS iot core will crash if this is not set.
_cleanSession :: Bool
, _lwt :: Maybe LastWill -- ^ LastWill message to be sent on client disconnect.
, _msgCB :: MessageCallback -- ^ Callback for incoming messages.
, _protocol :: ProtocolLevel -- ^ Protocol to use for the connection.
, _connProps :: [Property] -- ^ Properties to send to the broker in the CONNECT packet.
, _hostname :: String -- ^ Host to connect to (parsed from the URI)
, _port :: Int -- ^ Port number (parsed from the URI)
-- | Unique connection ID (parsed from the URI).
-- Set it by using x.y#connID
, _connID :: String
, _username :: Maybe String -- ^ Optional username (parsed from the URI)
, _password :: Maybe String -- ^ Optional password (parsed from the URI)
, _connectTimeout :: Int -- ^ Connection timeout (microseconds)
, _tlsSettings :: TLSSettings -- ^ TLS Settings for secure connections
, _pingPeriod :: Int -- ^ Time in seconds between pings
, _pingPatience :: Int -- ^ Time in seconds for which there must be no incoming packets before the broker is considered dead. Should be more than the ping period plus the maximum expected ping round trip time.
}
-- | A default 'MQTTConfig'. A '_connID' /may/ be required depending on
-- your broker (or if you just want an identifiable/resumable
-- connection). In MQTTv5, an empty connection ID may be sent and the
-- server may assign an identifier for you and return it in the
-- 'PropAssignedClientIdentifier' 'Property'.
mqttConfig :: MQTTConfig
mqttConfig = MQTTConfig{_hostname="", _port=1883, _connID="",
_username=Nothing, _password=Nothing,
_cleanSession=True, _lwt=Nothing,
_msgCB=NoCallback,
_protocol=Protocol311, _connProps=mempty,
_connectTimeout=180000000,
_tlsSettings=def,
_pingPeriod=30000000,
_pingPatience=90000000}
-- | Connect to an MQTT server by URI.
--
-- @mqtt://@, @mqtts://@, @ws://@, and @wss://@ URLs are supported.
-- The host, port, username, and password will be derived from the URI
-- and the values supplied in the config will be ignored.
--
-- > main :: IO
-- > main = do
-- > let (Just uri) = parseURI "mqtt://test.mosquitto.org"
-- > mc <- connectURI mqttConfig{} uri
-- > publish mc "tmp/topic" "hello!" False
connectURI :: MQTTConfig -> URI -> IO MQTTClient
connectURI cfg@MQTTConfig{..} uri = do
let cf = case uriScheme uri of
"mqtt:" -> runClient
"mqtts:" -> runClientTLS
"ws:" -> runWS uri False
"wss:" -> runWS uri True
us -> mqttFail $ "invalid URI scheme: " <> us
a = fromMaybe nullURIAuth $ uriAuthority uri
(u,p) = up (uriUserInfo a)
v <- namedTimeout "MQTT connect" _connectTimeout $
cf cfg{Network.MQTT.Client._connID=cid _protocol (uriFragment uri),
_hostname=uriRegName a, _port=port (uriPort a) (uriScheme uri),
Network.MQTT.Client._username=u, Network.MQTT.Client._password=p}
maybe (mqttFail $ "connection to " <> show uri <> " timed out") pure v
where
port "" "mqtt:" = 1883
port "" "mqtts:" = 8883
port "" "ws:" = 80
port "" "wss:" = 443
port x _ = (read . tail) x
cid _ ['#'] = ""
cid _ ('#':xs) = xs
cid _ _ = ""
up "" = (_username, _password)
up x = let (u,r) = break (== ':') (init x) in
(Just (unEscapeString u), if r == "" then Nothing else Just (unEscapeString $ tail r))
-- | Set up and run a client from the given config.
runClient :: MQTTConfig -> IO MQTTClient
runClient cfg@MQTTConfig{..} = tcpCompat (runTCPClient (clientSettings _port (BCS.pack _hostname))) cfg
-- | Set up and run a client connected via TLS.
runClientTLS :: MQTTConfig -> IO MQTTClient
runClientTLS cfg@MQTTConfig{..} = tcpCompat (runTLSClient tlsConf) cfg
where tlsConf = (tlsClientConfig _port (BCS.pack _hostname)) {tlsClientTLSSettings=_tlsSettings}
-- Compatibility mechanisms for TCP Conduit bits.
tcpCompat :: ((AppData -> IO ()) -> IO ()) -> MQTTConfig -> IO MQTTClient
tcpCompat mkconn = runMQTTConduit (adapt mkconn)
where adapt mk f = mk (f . adaptor)
adaptor ad = (appSource ad, appSink ad)
runWS :: URI -> Bool -> MQTTConfig -> IO MQTTClient
runWS URI{uriPath, uriQuery} secure cfg@MQTTConfig{..} =
runMQTTConduit (adapt $ cf secure _hostname _port endpoint WS.defaultConnectionOptions hdrs) cfg
where
hdrs = [("Sec-WebSocket-Protocol", "mqtt")]
adapt mk f = mk (f . adaptor)
adaptor s = (wsSource s, wsSink s)
-- An empty path would make a request line of "GET HTTP/1.1".
endpoint = (if null uriPath then "/" else uriPath) <> uriQuery
cf :: Bool -> String -> Int -> String -> WS.ConnectionOptions -> WS.Headers -> WS.ClientApp () -> IO ()
cf False = WS.runClientWith
cf True = runWSS
-- The connection closing ends the stream, as EOF does over TCP.
wsSource :: WS.Connection -> ConduitT () BCS.ByteString IO ()
wsSource ws = liftIO (E.try (WS.receiveData ws)) >>= \case
Left (WS.CloseRequest _ _) -> pure ()
Left WS.ConnectionClosed -> pure ()
Left e -> liftIO (E.throwIO e)
Right bs -> unless (BCS.null bs) (yield bs) *> wsSource ws
wsSink :: WS.Connection -> ConduitT BCS.ByteString Void IO ()
wsSink ws = maybe (pure ()) (\bs -> liftIO (WS.sendBinaryData ws bs) >> wsSink ws) =<< await
runWSS :: String -> Int -> String -> WS.ConnectionOptions -> WS.Headers -> WS.ClientApp () -> IO ()
runWSS host port path options hdrs' app = do
let connectionParams = ConnectionParams
{ connectionHostname = host
, connectionPort = toEnum port
, connectionUseSecure = Just _tlsSettings
, connectionUseSocks = Nothing
}
context <- initConnectionContext
E.bracket (connectTo context connectionParams) connectionClose
(\conn -> do
stream <- makeStream (reader conn) (writer conn)
-- Host carries the port, as it does for ws://.
WS.runClientWithStream stream (host <> ":" <> show port) path options hdrs' app)
where
reader conn =
catchIOError (Just <$> connectionGetChunk conn)
(\e -> if isEOFError e then pure Nothing else E.throwIO e)
writer conn = maybe (pure ()) (connectionPut conn . BC.toStrict)
mqttFail :: String -> a
mqttFail = E.throw . MQTTException
-- A couple async utilities to get our stuff named.
namedAsync :: String -> IO a -> IO (Async a)
namedAsync s a = async a >>= \p -> labelThread (asyncThreadId p) s >> pure p
namedTimeout :: String -> Int -> IO a -> IO (Maybe a)
namedTimeout n to a = timeout to (myThreadId >>= \tid -> labelThread tid n >> a)
-- | MQTTConduit provides a source and sink for data as used by 'runMQTTConduit'.
type MQTTConduit = (ConduitT () BCS.ByteString IO (), ConduitT BCS.ByteString Void IO ())
-- | Set up and run a client with a conduit context function.
--
-- The provided action calls another IO action with a 'MQTTConduit' as a
-- parameter. It is expected that this action will manage the
-- lifecycle of the conduit source/sink on behalf of the client.
runMQTTConduit :: ((MQTTConduit -> IO ()) -> IO ()) -- ^ an action providing an 'MQTTConduit' in an execution context
-> MQTTConfig -- ^ the 'MQTTConfig'
-> IO MQTTClient
runMQTTConduit mkconn MQTTConfig{..} = do
_ch <- newTChanIO
_pktID <- newTVarIO 1
_acks <- newTVarIO mempty
_inflight <- newTVarIO mempty
_st <- newTVarIO Starting
_ct <- newTVarIO Nothing
_outA <- newTVarIO mempty
_inA <- newTVarIO mempty
_connACKFlags <- newTVarIO (ConnACKFlags NewSession ConnUnspecifiedError mempty)
_corr <- newTVarIO mempty
_cbQ <- newTQueueIO
_cbHandle <- newTVarIO Nothing
_outFlight <- newTVarIO 0
let _cb = _msgCB
cli = MQTTClient{..}
t <- namedAsync "MQTT clientThread" $ clientThread cli
-- If we don't get connected (including a timeout or other async
-- exception, e.g. from connectURI), the client thread must not carry on
-- and connect anyway with nobody holding the client.
flip E.onException (cancel t) $ do
s <- atomically (waitForLaunch cli t)
when (s == Disconnected) $ wait t
atomically $ checkConnected cli
-- Ordered callbacks get messages from the start, not only after a
-- subscribe (a resumed session can deliver without one).
runCallbackThread cli
pure cli
where
clientThread cli = E.finally connectAndRun markDisco
where
connectAndRun = mkconn $ \ad -> start cli ad >>= run ad
markDisco = atomically $
modifyTVar' (_st cli) $ \case
Starting -> Disconnected
Connected -> Disconnected
Stopping -> Stopped
x -> x
start c (_,sink) = do
void . runConduit $ do
let req = connectRequest{T._connID=BC.pack _connID,
T._lastWill=_lwt,
T._username=TEL.encodeUtf8 . TL.pack <$> _username,
T._password=BC.pack <$> _password,
T._cleanSession=_cleanSession,
T._connProperties=_connProps}
yield (BL.toStrict $ toByteString _protocol req) .| sink
pure c
run (src,sink) c@MQTTClient{..} = do
pch <- newTChanIO
o <- namedAsync "MQTT out" $ onceConnected >> processOut
p <- namedAsync "MQTT ping" $ onceConnected >> doPing
w <- namedAsync "MQTT watchdog" $ watchdog pch
s <- namedAsync "MQTT in" $ doSrc pch
void $ waitAnyCancel [o, p, w, s]
where
doSrc pch = runConduit $ src
.| conduitParser (parsePacket _protocol)
.| C.mapM_ (\(_,x) -> liftIO (dispatch c pch $! x))
-- Past Starting, not exactly Connected: a disconnect can move us on
-- to Stopping before these threads first look.
onceConnected = atomically $ check . (/= Starting) =<< readTVar _st
processOut = runConduit $ outgoing .| C.map (BL.toStrict . toByteString _protocol) .| sink
-- Packets go out while connected, and while stopping (to get the
-- DISCONNECT out). After a DISCONNECT, the client closes the
-- connection itself [MQTT-3.14.4-1].
outgoing = do
p <- liftIO . atomically $ do
st <- readTVar _st
unless (st == Stopping) $ checkConnected c
readTChan _ch
yield p
case p of
DisconnectPkt _ -> pure ()
_ -> outgoing
-- Ping at least twice per Server Keep Alive, if the server set
-- one shorter than our period [MQTT-3.2.2-21].
doPing = do
ConnACKFlags _ _ props <- readTVarIO _connACKFlags
let period = minimum (_pingPeriod : [fromIntegral k * 500000 | PropServerKeepAlive k <- props, k > 0])
forever $ threadDelay period *> atomically (sendQuietly c PingPkt)
watchdog ch = forever $ do
toch <- registerDelay _pingPatience
timedOut <- atomically $ ((check =<< readTVar toch) >> pure True) `orElse` (readTChan ch >> pure False)
when timedOut $ killConn c Timeout
waitForLaunch MQTTClient{..} t = do
writeTVar _ct (Just t)
c <- readTVar _st
if c == Starting then retry else pure c
-- | Wait for a client to terminate its connection.
-- An exception is thrown if the client didn't terminate expectedly.
waitForClient :: MQTTClient -> IO ()
waitForClient c@MQTTClient{..} = do
r <- traverse waitCatch =<< readTVarIO _ct
st <- readTVarIO _st
case (st, r) of
-- A stopped client is fine however its thread ended (stopClient
-- cancels it).
(Stopped, _) -> pure ()
(_, Just (Left e)) | notCanceled e -> E.throwIO e
_ -> maybe (pure ()) E.throwIO =<< atomically (stateX c Stopped)
where notCanceled e = isNothing (E.fromException e :: Maybe AsyncCancelled)
-- | Stops the client and closes the connection without sending a DISCONNECT
-- message to the broker. This will cause the last-will message to be delivered
-- by the broker if it has been defined.
stopClient :: MQTTClient -> IO ()
stopClient MQTTClient{..} = do
-- Stopping first, so the client thread ends as Stopped rather than
-- Disconnected for anyone waiting on it.
atomically (writeTVar _st Stopping)
void . traverse cancel =<< readTVarIO _ct
void . traverse cancel =<< readTVarIO _cbHandle
atomically (writeTVar _st Stopped)
stateX :: MQTTClient -> ConnState -> STM (Maybe E.SomeException)
stateX MQTTClient{..} want = f <$> readTVar _st
where
je = Just . E.toException . MQTTException
f :: ConnState -> Maybe E.SomeException
f Connected = if want == Connected then Nothing else je "unexpectedly connected"
f Stopped = if want == Stopped then Nothing else je "unexpectedly stopped"
f Stopping = je "stopping"
f Disconnected = je "disconnected"
f Starting = je "died while starting"
f (DiscoErr x) = Just . E.toException . Discod $ x
f (ConnErr e) = je (show e)
data MQTTException = Timeout | BadData | Discod DisconnectRequest | MQTTException String deriving(Eq, Show)
instance E.Exception MQTTException
dispatch :: MQTTClient -> TChan Bool -> MQTTPkt -> IO ()
dispatch c@MQTTClient{..} pch pkt =
case pkt of
(ConnACKPkt p) -> connACKd p
(PublishPkt p) -> validTopic p >> pub p
(SubACKPkt (SubscribeResponse i _ _)) -> delegate DSubACK i
(UnsubACKPkt (UnsubscribeResponse i _ _)) -> delegate DUnsubACK i
(PubACKPkt (PubACK i _ _)) -> delegate DPubACK i
(PubRELPkt (PubREL i _ _)) -> pubd i
(PubRECPkt (PubREC i _ _)) -> delegate DPubREC i
(PubCOMPPkt (PubCOMP i _ _)) -> delegate DPubCOMP i
(DisconnectPkt req) -> disco req
PongPkt -> atomically . writeTChan pch $ True
-- Not implemented
(AuthPkt p) -> mqttFail ("unexpected incoming auth: " <> show p)
-- Things clients shouldn't see
PingPkt -> mqttFail "unexpected incoming ping packet"
(ConnPkt _ _) -> mqttFail "unexpected incoming connect"
(SubscribePkt _) -> mqttFail "unexpected incoming subscribe"
(UnsubscribePkt _) -> mqttFail "unexpected incoming unsubscribe"
where connACKd connr@(ConnACKFlags _ val _) = case val of
ConnAccepted -> atomically $ do
writeTVar _connACKFlags connr
writeTVar _st Connected
_ -> do
t <- readTVarIO _ct
atomically $ writeTVar _st (ConnErr connr)
maybeCancelWith (MQTTException $ show connr) t
-- A topic name can't have wildcards [MQTT-3.3.2-2]. One that does
-- is a protocol error ending the connection, rather than an
-- exception in whatever callback looks at the topic. (An empty one
-- is a topic alias.)
validTopic PublishRequest{_pubTopic}
| BL.null _pubTopic || isJust (mkTopic (blToText _pubTopic)) = pure ()
| otherwise = do
sendPacketIO c (DisconnectPkt (DisconnectRequest DiscoProtocolError mempty))
E.throwIO (MQTTException ("invalid topic in PUBLISH: " <> show _pubTopic))
pub p@PublishRequest{_pubQoS=QoS0} = atomically (resolve p) >>= notify (pure ())
pub p@PublishRequest{_pubQoS=QoS1, _pubPktID} =
notify (sendQuietly c (PubACKPkt (PubACK _pubPktID 0 mempty))) =<< atomically (resolve p)
-- QoS 2 is delivered on PUBLISH, since the PUBREC takes ownership of
-- the message [MQTT-4.3.3-9]: delivering on PUBREL would lose it if
-- the connection dropped in between. Only the packet ID is kept,
-- to drop duplicates until the PUBREL.
pub p@PublishRequest{_pubQoS=QoS2, _pubPktID=i} = do
(p', fresh) <- atomically $ do
p' <- resolve p
seen <- Map.member i <$> readTVar _inflight
unless seen $ modifyTVar' _inflight (Map.insert i Delivering)
sendQuietly c (PubRECPkt (PubREC i 0 mempty))
pure (p', not seen)
when fresh $ notify (qos2Done i) p'
qos2Done i = Map.lookup i <$> readTVar _inflight >>= \case
Just Delivering -> modifyTVar' _inflight (Map.insert i Delivered)
Just Released -> modifyTVar' _inflight (Map.delete i) *> pubcomp i
_ -> pure ()
-- A PUBREL for an ID we don't know was delivered already (e.g. on
-- an earlier connection of this session), so it's completed too.
pubd i = atomically $ Map.lookup i <$> readTVar _inflight >>= \case
Just Delivering -> modifyTVar' _inflight (Map.insert i Released)
Just Released -> pure ()
_ -> modifyTVar' _inflight (Map.delete i) *> pubcomp i
pubcomp i = sendQuietly c (PubCOMPPkt (PubCOMP i 0 mempty))
-- done runs once the callback has finished (sending the ack).
notify done p@PublishRequest{..} = do
cb <- maybe (pure _cb) (\cd -> Map.findWithDefault _cb cd <$> readTVarIO _corr) cdata
E.evaluate . force =<< case cb of
NoCallback -> pure ()
SimpleCallback f -> call (f c (blToTopic _pubTopic) _pubBody _pubProps)
OrderedCallback f -> callOrd (f c (blToTopic _pubTopic) _pubBody _pubProps)
LowLevelCallback f -> call (f c p)
OrderedLowLevelCallback f -> callOrd (f c p)
where
call a = link =<< namedAsync "notifier" (a *> atomically done)
callOrd a = atomically $ writeTQueue _cbQ (a *> atomically done)
cdata = foldr f Nothing _pubProps
where f (PropCorrelationData x) _ = Just x
f _ o = o
resolve p@PublishRequest{..} = do
topic <- resolveTopic (foldr aliasID Nothing _pubProps)
pure p{_pubTopic=topicToBL topic}
where
aliasID (PropTopicAlias x) _ = Just x
aliasID _ o = o
resolveTopic Nothing = pure (blToTopic _pubTopic)
resolveTopic (Just x) = do
when (_pubTopic /= "") $ modifyTVar' _inA (Map.insert x (blToTopic _pubTopic))
maybe (mqttFail ("failed to lookup topic alias " <> show x)) pure . Map.lookup x =<< readTVar _inA
delegate dt pid = atomically $
maybe (nak dt) (`writeTChan` pkt) . Map.lookup (dt, pid) =<< readTVar _acks
where
nak DPubREC = sendQuietly c (PubRELPkt (PubREL pid 0x92 mempty))
nak _ = pure ()
disco req = do
atomically $ writeTVar _st (DiscoErr req)
maybeCancelWith (Discod req) =<< readTVarIO _ct
-- Run OrderedCallbacks in a background thread. Does nothing for other callback types.
-- We keep the async handle in a TVar and make sure only one of these threads is running.
runCallbackThread :: MQTTClient -> IO ()
runCallbackThread c@MQTTClient{_cb, _cbQ, _cbHandle, _st}
| isOrdered _cb = do
-- We always spawn a thread, but may kill it if we already have
-- one. The new thread won't start until we atomically confirm
-- it's the only one.
latch <- newTVarIO False
handle <- namedAsync "ordered callbacks" (waitFor latch *> runOrderedCallbacks)
join . atomically $ do
cbThread <- readTVar _cbHandle
case cbThread of
Nothing -> do
-- This is the first thread. Flip the latch and put it in place.
writeTVar latch True
writeTVar _cbHandle (Just handle)
pure (pure ())
Just _ -> pure (cancel handle) -- otherwise, cancel the temporary thread.
| otherwise = pure ()
where isOrdered (OrderedCallback _) = True
isOrdered (OrderedLowLevelCallback _) = True
isOrdered _ = False
waitFor latch = atomically (check =<< readTVar latch)
-- Keep running callbacks from the queue. Once the connection
-- has ended, finish what's queued (messages already received),
-- then stop. A callback that fails ends the connection, as with
-- unordered callbacks, rather than leaving it up with nothing to
-- run the rest.
runOrderedCallbacks = loop `E.catch` \(e :: E.SomeException) ->
case E.fromException e of
Just (_ :: E.SomeAsyncException) -> E.throwIO e
Nothing -> killConn c e >> E.throwIO e
loop = join . atomically $ ((*> loop) <$> readTQueue _cbQ) `orElse` (readTVar _st >>= check . ended >> pure (pure ()))
ended st = st `notElem` [Starting, Connected, Stopping]
maybeCancelWith :: E.Exception e => e -> Maybe (Async ()) -> IO ()
maybeCancelWith e = void . traverse (`cancelWith` e)
killConn :: E.Exception e => MQTTClient -> e -> IO ()
killConn MQTTClient{..} e = readTVarIO _ct >>= maybeCancelWith e
checkConnected :: MQTTClient -> STM ()
checkConnected mc = maybe (pure ()) E.throw =<< stateX mc Connected
-- | True if we're currently in a normally connected state (in the IO monad).
isConnected :: MQTTClient -> IO Bool
isConnected = atomically . isConnectedSTM
-- | True if we're currently in a normally connected state (in the STM monad).
isConnectedSTM :: MQTTClient -> STM Bool
isConnectedSTM MQTTClient{..} = (Connected ==) <$> readTVar _st
sendPacket :: MQTTClient -> MQTTPkt -> STM ()
sendPacket c@MQTTClient{..} p = checkConnected c >> writeTChan _ch p
sendPacketIO :: MQTTClient -> MQTTPkt -> IO ()
sendPacketIO c = atomically . sendPacket c
-- | Queue a packet if the connection is still up, otherwise drop it. For
-- acks and pings, which shouldn't end (or race) a connection that is
-- already going away.
sendQuietly :: MQTTClient -> MQTTPkt -> STM ()
sendQuietly MQTTClient{..} p = readTVar _st >>= \st -> when (st == Connected || st == Stopping) (writeTChan _ch p)
textToBL :: Text -> BL.ByteString
textToBL = BL.fromStrict . TE.encodeUtf8
blToText :: BL.ByteString -> Text
blToText = TE.decodeUtf8 . BL.toStrict
topicToBL :: Topic -> BL.ByteString
topicToBL = textToBL . unTopic
filterToBL :: Filter -> BL.ByteString
filterToBL = textToBL . unFilter
blToTopic :: BL.ByteString -> Topic
blToTopic = fromJust . mkTopic . blToText
-- | Reserve the next packet ID that isn't still in use by an exchange
-- waiting for its response. If all of them are, wait for one to free up.
reservePktID :: MQTTClient -> [DispatchType] -> STM (TChan MQTTPkt, Word16)
reservePktID c@MQTTClient{..} dts = do
checkConnected c
ch <- newTChan
acks <- readTVar _acks
start <- readTVar _pktID
let next pid = if pid == maxBound then 1 else succ pid
inUse pid = any (\t -> Map.member (t, pid) acks) [minBound ..]
candidates = take (fromIntegral (maxBound :: Word16)) (iterate next start)
case filter (not . inUse) candidates of
[] -> retry
(pid:_) -> do
writeTVar _pktID (next pid)
writeTVar _acks (Map.union (Map.fromList [((t, pid), ch) | t <- dts]) acks)
pure (ch,pid)
-- Releasing an ID doesn't need a connection, and requiring one would
-- replace whatever exception is in flight with "disconnected".
releasePktID :: MQTTClient -> (DispatchType,Word16) -> STM ()
releasePktID MQTTClient{..} k = modifyTVar' _acks (Map.delete k)
releasePktIDs :: MQTTClient -> [(DispatchType,Word16)] -> STM ()
releasePktIDs MQTTClient{..} ks = modifyTVar' _acks deleteMany
where deleteMany m = Map.withoutKeys m (Set.fromList ks)
sendAndWait :: MQTTClient -> DispatchType -> (Word16 -> MQTTPkt) -> IO MQTTPkt
sendAndWait c@MQTTClient{..} dt f = do
(ch,pid) <- atomically $ do
(ch,pid) <- reservePktID c [dt]
sendPacket c (f pid)
pure (ch,pid)
-- Wait for the response in a separate transaction.
atomically $ do
st <- readTVar _st
when (st /= Connected) $ mqttFail "disconnected waiting for response"
releasePktID c (dt,pid)
readTChan ch
-- | Subscribe to a list of topic filters with their respective 'QoS'es.
-- The accepted 'QoS'es are returned in the same order as requested.
subscribe :: MQTTClient -> [(Filter, SubOptions)] -> [Property] -> IO ([Either SubErr QoS], [Property])
subscribe c ls props = do
runCallbackThread c
sendAndWait c DSubACK (\pid -> SubscribePkt $ SubscribeRequest pid ls' props) >>= \case
SubACKPkt (SubscribeResponse _ rs aprops) -> pure (rs, aprops)
pkt -> mqttFail $ "unexpected response to subscribe: " <> show pkt
where ls' = map (first filterToBL) ls
-- | Unsubscribe from a list of topic filters.
--
-- In MQTT 3.1.1, there is no body to an unsubscribe response, so it
-- can be ignored. If this returns, you were unsubscribed. In MQTT
-- 5, you'll get a list of unsub status values corresponding to your
-- request filters, and whatever properties the server thought you
-- should know about.
unsubscribe :: MQTTClient -> [Filter] -> [Property] -> IO ([UnsubStatus], [Property])
unsubscribe c ls props = do
(UnsubACKPkt (UnsubscribeResponse _ rsn rprop)) <- sendAndWait c DUnsubACK (\pid -> UnsubscribePkt $ UnsubscribeRequest pid (map filterToBL ls) props)
pure (rprop, rsn)
-- | Publish a message (QoS 0).
publish :: MQTTClient
-> Topic -- ^ Topic
-> BL.ByteString -- ^ Message body
-> Bool -- ^ Retain flag
-> IO ()
publish c t m r = void $ publishq c t m r QoS0 mempty
-- | Publish a message with the specified QoS and Properties list.
publishq :: MQTTClient
-> Topic -- ^ Topic
-> BL.ByteString -- ^ Message body
-> Bool -- ^ Retain flag
-> QoS -- ^ QoS
-> [Property] -- ^ Properties
-> IO ()
publishq c t m r q props = publishWith c (pure (t, props)) m r q
-- | Publish with the topic and properties from an STM action. The action
-- runs in the same transaction that queues the PUBLISH, so whatever it
-- decides (e.g. a topic alias) goes out in the order it was decided.
publishWith :: MQTTClient -> STM (Topic, [Property]) -> BL.ByteString -> Bool -> QoS -> IO ()
publishWith c@MQTTClient{_outFlight, _connACKFlags} mk m r q
-- QoS 0 has no packet ID and nothing to wait for.
| q == QoS0 = atomically $ do
(t, props) <- mk
sendPacket c (pkt t props 0)
| otherwise = do
(ch,pid) <- atomically $ do
(ch,pid) <- reservePktID c types
takeSlot
(t, props) <- mk
sendPacket c (pkt t props pid)
pure (ch,pid)
E.finally (satisfyQoS ch pid) . atomically $ do
modifyTVar' _outFlight (subtract 1)
releasePktIDs c [(t',pid) | t' <- types]
where
types = [DPubACK, DPubREC, DPubCOMP]
-- No more QoS 1 and 2 PUBLISHes in flight than the server's Receive
-- Maximum [MQTT-3.3.4-7].
takeSlot = do
ConnACKFlags _ _ props <- readTVar _connACKFlags
n <- readTVar _outFlight
check (n < foldr (\p o -> case p of PropReceiveMaximum x -> x; _ -> o) maxBound props)
writeTVar _outFlight (n + 1)
pkt t props pid = PublishPkt $ PublishRequest {_pubDup = False,
_pubQoS = q,
_pubPktID = pid,
_pubRetain = r,
_pubTopic = topicToBL t,
_pubBody = m,
_pubProps = props}
satisfyQoS ch pid
| q == QoS0 = pure ()
| q == QoS1 = void $ do
(PubACKPkt (PubACK _ st pprops)) <- atomically $ checkConnected c >> readTChan ch
unless (isOK st) $ mqttFail ("qos 1 publish error: " <> show st <> " " <> show pprops)
| q == QoS2 = waitRec
| otherwise = error "invalid QoS"
where
isOK 0 = True -- success
isOK 16 = True -- It worked, but nobody cares (no matching subscribers)
isOK _ = False
-- The packet ID is held until the PUBCOMP, since until then
-- the server still has the exchange open [MQTT-4.3.3].
waitRec = atomically (checkConnected c >> readTChan ch) >>= \case
PubRECPkt (PubREC _ st recprops) -> do
unless (isOK st) $ mqttFail ("qos 2 REC publish error: " <> show st <> " " <> show recprops)
sendPacketIO c (PubRELPkt $ PubREL pid 0 mempty)
waitRec
PubCOMPPkt (PubCOMP _ st' compprops) ->
when (st' /= 0) $ mqttFail ("qos 2 COMP publish error: " <> show st' <> " " <> show compprops)
wtf -> mqttFail ("unexpected packet received in QoS2 publish: " <> show wtf)
-- | Disconnect from the MQTT server.
disconnect :: MQTTClient -> DiscoReason -> [Property] -> IO ()
disconnect c@MQTTClient{..} reason props = race_ getDisconnected orDieTrying
where
-- Stopping in the same transaction that queues the DISCONNECT, so the
-- connection ending (from either side) after it counts as a stop.
getDisconnected = do
atomically $ do
checkConnected c
writeTVar _st Stopping
writeTChan _ch (DisconnectPkt $ DisconnectRequest reason props)
void . traverse waitCatch =<< readTVarIO _ct
atomically $ writeTVar _st Stopped
orDieTrying = threadDelay 10000000 *> atomically (writeTVar _st Disconnected) *> killConn c Timeout
-- | Disconnect with 'DiscoNormalDisconnection' and no properties.
normalDisconnect :: MQTTClient -> IO ()
normalDisconnect c = disconnect c DiscoNormalDisconnection mempty
-- | A convenience method for creating a 'LastWill'.
mkLWT :: Topic -> BL.ByteString -> Bool -> T.LastWill
mkLWT t m r = T.LastWill{
T._willRetain=r,
T._willQoS=QoS0,
T._willTopic = topicToBL t,
T._willMsg=m,
T._willProps=mempty
}
-- | Get the list of properties that were sent from the broker at connect time.
svrProps :: MQTTClient -> IO [Property]
svrProps = fmap p . atomically . connACKSTM
where p (ConnACKFlags _ _ props) = props
-- | Get the complete connection ACK packet from the beginning of this session.
connACKSTM :: MQTTClient -> STM ConnACKFlags
connACKSTM MQTTClient{_connACKFlags} = readTVar _connACKFlags
-- | Get the complete connection aCK packet from the beginning of this session.
connACK :: MQTTClient -> IO ConnACKFlags
connACK = atomically . connACKSTM
maxAliases :: MQTTClient -> IO Word16
maxAliases = fmap (foldr f 0) . svrProps
where
f (PropTopicAliasMaximum n) _ = n
f _ o = o
-- | Publish a message with the specified 'QoS' and 'Property' list. If
-- possible, use an alias to shorten the message length. The alias
-- list is managed by the client in a first-come, first-served basis,
-- so if you use this with more properties than the broker allows,
-- only the first N (up to TopicAliasMaximum, as specified by the
-- broker at connect time) will be aliased.
--
-- This is safe to use as a general publish mechanism, as it will
-- default to not aliasing whenver there's not already an alias and we
-- can't create any more.
pubAliased :: MQTTClient
-> Topic -- ^ Topic
-> BL.ByteString -- ^ Message body
-> Bool -- ^ Retain flag
-> QoS -- ^ QoS
-> [Property] -- ^ Properties
-> IO ()
pubAliased c@MQTTClient{..} t m r q props = do
x <- maxAliases c
-- The alias is chosen in the transaction that queues the PUBLISH, so
-- the one that introduces it can't be overtaken by one that uses it.
publishWith c (alias x) m r q
where
alias mv = do
as <- readTVar _outA
let n = toEnum (length as + 1)
cur = Map.lookup t as
v = fromMaybe (if n > mv then 0 else n) cur
when (v > 0) $ writeTVar _outA (Map.insert t v as)
pure (maybe t (const "") cur, props <> [PropTopicAlias v | v > 0])
-- | Register a callback handler for a message with the given correlated data identifier.
--
-- This registration will remain in place until unregisterCorrelated is called to remove it.
registerCorrelated :: MQTTClient -> BL.ByteString -> MessageCallback -> STM ()
registerCorrelated MQTTClient{_corr} bs cb = modifyTVar' _corr (Map.insert bs cb)
-- | Unregister a callback handler for the given correlated data identifier.
unregisterCorrelated :: MQTTClient -> BL.ByteString -> STM ()
unregisterCorrelated MQTTClient{_corr} bs = modifyTVar' _corr (Map.delete bs)