diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c)2011, Vitaliy Rukavishnikov + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of the author nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README b/README new file mode 100644 --- /dev/null +++ b/README @@ -0,0 +1,7 @@ +Stomp is a client library for communicating with message servers through the +STOMP protocol (http://stomp.github.com/stomp-specification-1.1.html) + + + + + 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/Stomp.cabal b/Stomp.cabal new file mode 100644 --- /dev/null +++ b/Stomp.cabal @@ -0,0 +1,28 @@ +Name: Stomp +Version: 0.1 +Description: Stomp is a client library for communicating with message servers through the + STOMP protocol (http://stomp.github.com/stomp-specification-1.1.html) +License: BSD3 +License-File: LICENSE +Author: Vitaliy Rukavishnikov +Maintainer: virukav@gmail.com +Homepage: http://github.com/rukav/Stomp +Bug-Reports: mailto:virukav@gmail.com +Build-Type: Simple +Category: Network +Synopsis: Client library for Stomp brokers. +Cabal-Version: >=1.2 +Extra-Source-Files: README + test/Tests.hs + examples/Example.hs +Library + Hs-Source-Dirs: src + Exposed-Modules: Network.Stomp + Build-Depends: base >= 4 && < 5, + network >= 2.3.0.2, + binary, + time, + bytestring >= 0.9, + utf8-string >= 0.3.6 + + diff --git a/examples/Example.hs b/examples/Example.hs new file mode 100644 --- /dev/null +++ b/examples/Example.hs @@ -0,0 +1,31 @@ +import Network.Stomp +import qualified Data.ByteString.Lazy.Char8 as B + +main = do + -- connect to a stomp broker + con <- connect "stomp://guest:guest@127.0.0.1:61613" vers headers + putStrLn $ "Accepted versions: " ++ show (versions con) + + -- start consumer and subscribe to the queue + startConsumer con callback + subscribe con "/queue/test" "0" [] + + -- send the messages to the queue + send con "/queue/test" [] (B.pack "message1") + send con "/queue/test" [] (B.pack "message2") + + -- wait + getLine + + -- unsubscribe and disconnect + unsubscribe con "0" [] + disconnect con [] + where + vers = [(1,0),(1,1)] + headers = [] + +callback :: Frame -> IO () +callback (Frame (SC MESSAGE) hs body) = do + putStrLn $ "received message: " ++ (B.unpack body) + putStrLn $ "headers: " ++ show hs +callback f = putStrLn $ "received frame: " ++ show f diff --git a/src/Network/Stomp.hs b/src/Network/Stomp.hs new file mode 100644 --- /dev/null +++ b/src/Network/Stomp.hs @@ -0,0 +1,595 @@ +{-# LANGUAGE ScopedTypeVariables, DeriveDataTypeable #-} +{- | + +A client library for Stomp serevers implementing stomp 1.1 specification. See http://stomp.github.com/stomp-specification-1.1.html + +/Example/: + +>import Network.Stomp +>import qualified Data.ByteString.Lazy.Char8 as B +> +>main = do +> -- connect to a stomp broker +> con <- connect "stomp://guest:guest@127.0.0.1:61613" vers headers +> putStrLn $ "Accepted versions: " ++ show (versions con) +> +> -- start consumer and subscribe to the queue +> startConsumer con callback +> subscribe con "/queue/test" "0" [] +> +> -- send the messages to the queue +> send con "/queue/test" [] (B.pack "message1") +> send con "/queue/test" [] (B.pack "message2") +> +> -- wait +> getLine +> +> -- unsubscribe and disconnect +> unsubscribe con "0" [] +> disconnect con [] +> where +> vers = [(1,0),(1,1)] +> headers = [] +> +>callback :: Frame -> IO () +>callback (Frame (SC MESSAGE) hs body) = do +> putStrLn $ "received message: " ++ (B.unpack body) +> putStrLn $ "headers: " ++ show hs +>callback f = putStrLn $ "received frame: " ++ show f + +-} + +module Network.Stomp ( + Command (..), + ClientCommand (..), + ServerCommand (..), + Frame (..), + Connection, + StompUri, + Host, + Destination, + MessageId, + Transaction, + Subscription, + Version, + StompException (..), + + connect, + stomp, + connect', + stomp', + disconnect, + send, + subscribe, + unsubscribe, + ack, + nack, + begin, + commit, + abort, + + startConsumer, + receiveFrame, + setExcpHandler, + startSendBeat, + startRecvBeat, + sendTimeout, + recvTimeout, + lastSend, + lastRecv, + versions +) +where + +import qualified Data.ByteString.UTF8 as BU +import qualified Data.ByteString.Char8 as BS +import qualified Data.ByteString.Lazy.Char8 as BL +import Data.List (intercalate) +import Data.Time (UTCTime, getCurrentTime, diffUTCTime) +import Data.Typeable +import Data.Maybe +import Data.Char (toLower) +import qualified Data.Binary.Builder as B +import Data.Binary.Put + +import Network.BSD +import Network.URI +import qualified Network.Socket as N + +import qualified Control.Exception as E +import Control.Concurrent +import Control.Monad +import System.IO + +-- | Stomp frame record +data Frame = Frame { + command :: Command, + headers :: [Header], + body :: BL.ByteString +} deriving Show + +-- | Client frame commands +data ClientCommand + = SEND + | SUBSCRIBE + | UNSUBSCRIBE + | BEGIN + | COMMIT + | ABORT + | ACK + | NACK + | DISCONNECT + | CONNECT + | STOMP + deriving (Show, Read, Eq) + +-- | Broker frame commands +data ServerCommand + = CONNECTED + | MESSAGE + | RECEIPT + | ERROR + deriving (Show, Read, Eq) + +-- | Stomp frame commands +data Command + = CC ClientCommand + | SC ServerCommand + deriving (Show, Read, Eq) + +-- | Stomp header as a name/value pair +type Header = (String, String) + +-- | A record used to communicate with Stomp brokers +data Connection = Connection { + handle :: Handle, + versions :: [Version], -- ^ accepted stomp versions + listener :: MVar ThreadId, -- ^ frames consumer thread + sendBeat :: MVar ThreadId, -- ^ send heartbeat thread + recvBeat :: MVar ThreadId, -- ^ recieve heartbeat thread + sendTimeout :: Int, -- ^ send heartbeat timeout + recvTimeout :: Int, -- ^ receive heaertbeat timeout + lastSend :: MVar UTCTime, -- ^ last frame sent time + lastRecv :: MVar UTCTime, -- ^ last frame received time + sockLock :: MVar (), -- ^ used for atomic writes to the socket + closed :: MVar (Maybe StompException), -- ^ closed connection flag + disconReq :: MVar String, -- ^ disconnect receipt request + disconResp :: MVar (), -- ^ disconnect receipt response + excpHandle :: MVar (Maybe (StompException -> IO ())) -- ^ exception handler in the frames consumer +} + +data StompException + = ConnectionError String + | InvalidUri String + | InvalidFrame String + | BrokerError String + | StompIOError E.IOException + deriving (Typeable, Show, Eq) + +instance E.Exception StompException + +type StompUri = String +type Host = String +type Destination = String +type MessageId = String +type Transaction = String +type Subscription = String +type Version = (Int,Int) + +------- Stomp API + +-- | connect to the stomp (1.0, 1.1) broker using uri +connect :: StompUri -> [Version] -> [Header] -> IO Connection +connect uri vs hs = do + (host, port, hs') <- processUri uri + connect' host port vs (hs ++ hs') + +-- | connect to the stomp 1.1 broker using uri +stomp :: StompUri -> [Header] -> IO Connection +stomp uri hs = do + (host, port, hs') <- processUri uri + stomp' host port (hs ++ hs') + +-- | connect to the stomp (1.0, 1.1) broker using hostname and port +connect' :: Host -> PortNumber -> [Version] -> [Header] -> IO Connection +connect' h p [] hs = mkConnection CONNECT h p (hs ++ [("accept-version","1.0")]) +connect' h p vs hs = mkConnection CONNECT h p hs' + where hs' = hs ++ [("accept-version", vers)] + vers = intercalate "," $ map go vs + go (v,v') = show v ++ '.':show v' + +-- | connect to the stomp 1.0 broker using hostname and port +stomp' :: Host -> PortNumber -> [Header] -> IO Connection +stomp' = mkConnection STOMP + +-- | closes stomp connection +disconnect :: Connection -> [Header] -> IO () +disconnect con hs = do + let receipt = lookup "receipt" hs + case receipt of + Nothing -> do stop con + sendFrame con (Frame (CC DISCONNECT) hs BL.empty) + Just r -> do putMVar (disconReq con) r + sendFrame con (Frame (CC DISCONNECT) hs BL.empty) + readMVar $ disconResp con + stop con + hClose $ handle con + modifyMVar_ (closed con) $ \x -> + return (Just $ fromMaybe (ConnectionError "Connection closed") x) + +{- send message to the destination. +The size of the message should be equal to the value of 'content-length' header. +-} +send :: Connection -> Destination -> [Header] -> BL.ByteString -> IO () +send con dest hs body = sendFrame con (Frame (CC SEND) hs' body) + where hs' = hs ++ [("destination", dest)] ++ clh + clh = maybe hdr (\_->[]) $ lookup "content-length" hs + hdr = [("content-length", show $ BL.length body)] + +-- | subscribe to the destination to receive stomp frames +subscribe :: Connection -> Destination -> Subscription -> [Header] -> IO () +subscribe con dest id hs = sendFrame con (Frame (CC SUBSCRIBE) hs' BL.empty) + where hs' = hs ++ [("destination", dest), ("id", id)] + +-- | unsubscribe from destination given the subscription id +unsubscribe :: Connection -> Subscription -> [Header] -> IO () +unsubscribe con id hs = sendFrame con (Frame (CC UNSUBSCRIBE) hs' BL.empty) + where hs' = hs ++ [("id", id)] + +-- | acknowledge the consumption of a message from a subscription +ack :: Connection -> Subscription -> MessageId -> [Header] -> IO () +ack con id msgid hs = sendFrame con (Frame (CC ACK) hs' BL.empty) + where hs' = hs ++ [("subscription", id), ("message-id", msgid)] + +-- | acknowledge the rejection of a message from a subscription +nack :: Connection -> Subscription -> MessageId -> [Header] -> IO () +nack con id msgid hs = sendFrame con (Frame (CC NACK) hs' BL.empty) + where hs' = hs ++ [("subscription", id), ("message-id", msgid)] + +-- | start a transaction +begin :: Connection -> Transaction -> [Header] -> IO () +begin con tid hs = sendFrame con (Frame (CC BEGIN) hs' BL.empty) + where hs' = hs ++ [("transaction", tid)] + +-- | commit a transaction +commit :: Connection -> Transaction -> [Header] -> IO () +commit con tid hs = sendFrame con (Frame (CC COMMIT) hs' BL.empty) + where hs' = hs ++ [("transaction", tid)] + +-- | rollback a transaction +abort :: Connection -> Transaction -> [Header] -> IO () +abort con tid hs = sendFrame con (Frame (CC ABORT) hs' BL.empty) + where hs' = hs ++ [("transaction", tid)] + +--------- Connection + +-- | opens connection to the stomp boker +mkConnection :: ClientCommand -> Host -> PortNumber -> [Header] -> IO Connection +mkConnection cmd host port hs = do + con <- newConn host port hs + sendFrame con (Frame (CC cmd) hs BL.empty) + (Frame (SC cmd) hs' body) <- receiveFrame con + when (cmd /= CONNECTED) $ + E.throwIO $ ConnectionError (BL.unpack body) + let (sendBeat, recvBeat) = getBeats hs hs' + return con {recvTimeout = recvBeat, sendTimeout = sendBeat, versions = ver hs'} + where ver h = maybe [(1,0)] parseVer $ lookup "version" h + +-- | parse versions supported by the broker +parseVer :: String -> [Version] +parseVer vs + | null r = [mkV v] + | otherwise = mkV v : parseVer (tail r) + where (v,r) = break (==',') vs + mkV s = let (m, m') = break (=='.') s in + (read m, read (tail m')) + +-- | open the server socket and convert it to handle +openSocket :: String -> PortNumber -> IO Handle +openSocket host port = do + proto <- getProtocolNumber "tcp" + sock <- N.socket N.AF_INET N.Stream proto + addr <- N.inet_addr host + N.connect sock (N.SockAddrInet port addr) + h <- N.socketToHandle sock ReadWriteMode + hSetBuffering h (BlockBuffering Nothing) + return h + +-- | initialize the new connection +newConn :: Host -> PortNumber -> [Header] -> IO Connection +newConn host port hs = do + h <- openSocket host port + lstnr <- newEmptyMVar + sBeat <- newEmptyMVar + rBeat <- newEmptyMVar + lSend <- newEmptyMVar + lRecv <- newEmptyMVar + sLock <- newMVar () + close <- newMVar Nothing + dRcpt <- newEmptyMVar + dLock <- newEmptyMVar + eHndl <- newMVar Nothing + return $ Connection h [] lstnr sBeat rBeat 0 0 lSend lRecv sLock close dRcpt dLock eHndl + +-- | create the authority value +stompAuth :: String -> Maybe URIAuth +stompAuth str = maybe Nothing auth (parseURI str) + where auth (URI s a _ _ _) = + if map toLower s /= "stomp:" then Nothing else a + +-- | retrieve host, port and headers from the authority value +fromAuth :: PortNumber -> URIAuth -> (Host, PortNumber, [Header]) +fromAuth defPort ua = + (host, portNum, [("login", user), ("passcode", pwd')]) + where (user, pwd) = break (==':') (uriUserInfo ua) + pwd' = if null pwd then [] else drop 1 $ init pwd + port = drop 1 (uriPort ua) + portNum = if null port then defPort + else toEnum (read port :: Int) + host = uriRegName ua + +-- | parse stomp uri +processUri :: String -> IO (Host, PortNumber, [Header]) +processUri uri = + maybe (E.throwIO $ InvalidUri uri) + (return . fromAuth defaultPort) + (stompAuth uri) + +-- | stop all children threads +stop :: Connection -> IO () +stop c = mapM_ go [listener c, sendBeat c, recvBeat c] + where go x = do + t <- tryTakeMVar x + maybe (return ()) killThread t + +defaultPort :: PortNumber +defaultPort = 61613 + +-------- Exceptions + +-- | set exception handler callback to process the exception in the consumer/heartbeats threads +setExcpHandler :: Connection -> (StompException -> IO ()) -> IO () +setExcpHandler con fun = + modifyMVar_ (excpHandle con) $ \_ -> return $ Just fun + +-- | invoke the exception handler and close the connection +onException :: Connection -> StompException -> IO () +onException con e = + withMVar (excpHandle con) $ \h -> do + maybe (return ()) (\f -> f e) h + modifyMVar_ (closed con) $ \x -> + return $ mplus x (Just e) + +-------- Consume frames + +-- | create consume frames thread +startConsumer :: Connection -> (Frame -> IO ()) -> IO () +startConsumer c fun = do + t <- tryTakeMVar (listener c) + when (isNothing t) $ do + tid <- forkIO $ E.finally + (consumeFrames c fun) + (tryPutMVar (disconResp c) ()) + tryPutMVar (listener c) tid + return () + +-- | For any incoming frame the user callback will be invoked +consumeFrames :: Connection -> (Frame -> IO ()) -> IO () +consumeFrames con fun = + E.catch + (do + frame <- receiveFrame con + fun frame + rec <- tryTakeMVar (disconReq con) + unless (isJust rec && checkReceipt (fromJust rec) frame) $ do + maybe (return ()) (putMVar (disconReq con)) rec + consumeFrames con fun) + (\(e::StompException) -> onException con e) + +-- | check receipt value in the frame headers +checkReceipt :: String -> Frame -> Bool +checkReceipt rec (Frame (SC RECEIPT) hs _) = + maybe False (== rec) $ lookup "receipt-id" hs +checkReceipt _ _ = False + +-- | send frame buffer to the server +sendFrame :: Connection -> Frame -> IO () +sendFrame con f = sendBuf con (strict $ runPut $ putFrame f) + where + putFrame frame@(Frame (CC cmd) hs body) = do + mapM_ (putByteString . BU.fromString) + [show cmd, "\n", hdrToStr cmd hs, "\n"] + unless (BL.null body) $ + putLazyByteString body + putWord8 0x00 + strict x = BS.concat (BL.toChunks x) + +{- unparse frame headers. +All frames except the CONNECT and CONNECTED frames will also escape any +colon or newline octets +-} +hdrToStr :: ClientCommand -> [Header] -> String +hdrToStr _ [] = [] +hdrToStr CONNECT hs = hdrToStr' id hs +hdrToStr cmd hs = hdrToStr' esc hs + +hdrToStr' f xs = unlines $ map go xs + where go (x,y) = f x ++ ":" ++ f y + +-- | send buffer to the handle +sendBuf :: Connection -> BS.ByteString -> IO () +sendBuf con bs = + withMVar (closed con) $ \c -> + if isJust c then E.throwIO (fromJust c) + else + E.catch + (withMVar (sockLock con) $ \_ -> do + BS.hPut (handle con) bs + hFlush (handle con) + beatTime (lastSend con)) + (\(e :: E.IOException) -> E.throwIO $ StompIOError e) + +--------- Receive frame + +-- | reads incoming frame from handle +receiveFrame :: Connection -> IO Frame +receiveFrame con = do + cmd <- readCommand con + headers <- readHeaders con cmd + body <- readBody con headers + beatTime (lastRecv con) + return (Frame cmd headers body) + +-- | read line from handle +readLine :: Handle -> IO String +readLine h = fmap BU.toString (BS.hGetLine h) + +-- | read stomp command or broker heartbeat from handle +readCommand :: Connection -> IO Command +readCommand con = do + eof <- hIsEOF (handle con) + if eof then + E.throwIO $ ConnectionError "Connection closed by broker" + else do + l <- readLine (handle con) + let l' = dropWhile (=='\x00') l + if null l' then do -- broker heartbeat + beatTime (lastRecv con) + readCommand con + else return $ SC (read l' :: ServerCommand) + +-- | read stomp headers from handle +readHeaders :: Connection -> Command -> IO [Header] +readHeaders con cmd = do + l <- readLine (handle con) + if null l then return [] + else do hs <- readHeaders con cmd + return (header (unesc l):hs) + case cmd of + (SC CONNECTED) -> return (header l:hs) + _ -> return (header (unesc l):hs) + where + header x = let (name, val) = break (==':') x + in (name, tail val) + +-- read stomp message with length 'content-length' or until frame terminator +readBody :: Connection -> [Header] -> IO BL.ByteString +readBody con hs = maybe (readTill h) (readBuf h) len + where + h = handle con + len = lookup "content-length" hs + readBuf h x = do + bs <- BL.hGet h (read x :: Int) + tr <- BL.hGet h 1 + when (tr /= term) $ + E.throwIO $ InvalidFrame "Missing frame terminator" + return bs + readTill h = liftM B.toLazyByteString (readTill' h) + readTill' h = do + ch <- BL.hGet h 1 + if ch == term then return B.empty + else liftM (B.append $ B.fromLazyByteString ch) (readTill' h) + +-- | stomp frame terminator +term :: BL.ByteString +term = BL.pack "\x00" + +-------- Escaping + +-- | escape header name and value +esc :: String -> String +esc = concatMap (\c -> fromMaybe [c] (code c)) + where code x = lookup x escMap + escMap = zip ":\n\\" ["\\c","\\n","\\\\"] + +-- | unescape header name and value +unesc :: String -> String +unesc [] = [] +unesc [x] = [x] +unesc (x:x':xs) + | [x,x'] == "\\n" = '\n' : unesc xs + | [x,x'] == "\\c" = ':' : unesc xs + | [x,x'] == "\\\\" = '\\' : unesc xs + | otherwise = x:unesc(x':xs) + +-------- Heartbeat + +{- get heartbeats values from the client and broker headers +The receiver tolerance is 100%. +-} +getBeats :: [Header] -> [Header] -> (Int, Int) +getBeats xs ys = (getBeat cs sr, 2 * getBeat cr ss) + where + getBeat x y + | x /= 0 && y /= 0 = max x y + | otherwise = 0 + (cs,cr) = beat xs + (ss,sr) = beat ys + +-- | parse heartbeat +beat :: [Header] -> (Int,Int) +beat = maybe (0,0) parseBeat . lookup "heart-beat" + where + parseBeat = parse . break (==',') + parse (x,y) = (read x :: Int, read (tail y) :: Int) + +-- | update heartbeat time +beatTime :: MVar UTCTime -> IO () +beatTime v = do + now <- getCurrentTime + b <- tryPutMVar v now + unless b $ modifyMVar_ v $ \_ -> return now + +-- | periodically send heartbeat data to the server +clientBeat :: Connection -> IO () +clientBeat con = + E.catch + (do + prev <- readMVar (lastSend con) + now <- getCurrentTime + let diff = floor $ diffUTCTime now prev * 1000 + let delay = sendTimeout con + if diff >= delay then do + sendBuf con (BU.fromString "\n\x00") + threadDelay (1000 * delay) + else threadDelay (1000 * (delay - diff)) + clientBeat con) + (\(e::StompException) -> onException con e) + +-- | periodically check a last frame receive time and mark connection as dead if needed +brokerBeat :: Connection -> IO () +brokerBeat con = + E.catch + (do + prev <- readMVar (lastRecv con) + now <- getCurrentTime + let diff = floor $ diffUTCTime now prev * 1000 + let delay = recvTimeout con + if diff >= delay then + E.throwIO $ BrokerError "Broker timeout expired" + else threadDelay (1000 * (delay - diff)) + brokerBeat con) + (\(e::StompException) -> onException con e) + +-- | fork send heartbeat thread +startSendBeat :: Connection -> IO () +startSendBeat c = do + svar <- tryTakeMVar (sendBeat c) + when (sendTimeout c > 0 && isNothing svar) $ do + tid <- forkIO $ E.finally + (clientBeat c) + (tryPutMVar (disconResp c) ()) + tryPutMVar (sendBeat c) tid + return () + +-- | fork receive heartbeat thread +startRecvBeat :: Connection -> IO () +startRecvBeat c = do + rvar <- tryTakeMVar (recvBeat c) + when (recvTimeout c > 0 && isNothing rvar) $ do + tid <- forkIO $ E.finally + (brokerBeat c) + (tryPutMVar (disconResp c) ()) + tryPutMVar (recvBeat c) tid + return () + diff --git a/test/Tests.hs b/test/Tests.hs new file mode 100644 --- /dev/null +++ b/test/Tests.hs @@ -0,0 +1,364 @@ +{-# LANGUAGE ScopedTypeVariables #-} +import Test.Framework (defaultMain, testGroup) +import Test.Framework.Providers.HUnit +import Test.HUnit + +import Network.Stomp +import qualified Data.ByteString.Char8 as BS +import qualified Data.ByteString.Lazy.Char8 as BL +import qualified Control.Exception as E +import Control.Concurrent +import Data.IORef +import Data.Maybe +import Control.Monad +import System.Timeout + +main = defaultMain tests +tests = [ + testGroup "Stomp commands" + [ testCase "connect" testConnect + , testCase "stomp" testStomp + , testCase "disconnect" testDisconnect + , testCase "send" testSend + , testCase "subscribe" testSubscribe + , testCase "unsubscribe" testUnsubscribe + , testCase "ack" testAck + , testCase "nack" testNack + , testCase "commit" testCommit + , testCase "abort" testAbort + ], + testGroup "Stomp protocol" + [ testCase "receipt" testReceipt + , testCase "sendbeat" testSendbeat + , testCase "recvbeat" testRecvbeat + , testCase "errors" testErrors + , testCase "encode" testEncode + ], + testGroup "Stomp messages" + [ testCase "null" testNull + , testCase "length" testLength + , testCase "big" testBig + ] + ] + +testConnect = do + (con, resp, _) <- mkStomp + E.finally + (do r <- readChan' resp + assertBool "Unexpected stomp response" (isNothing r)) + (disconnect con []) + +testStomp = do + con <- connect uri [] [] + disconnect con [] + unless (maximum (versions con) == (1,0)) $ do + con' <- stomp uri [] + disconnect con' [] + +testDisconnect = do + (con, resp, _) <- mkStomp + disconnect con [] + E.catch + (do disconnect con [] + assertBool "Unexpected valid connection" False) + (\(e::StompException) -> + assertBool "Invalid exception" (isConError e)) + +testSend = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do send con dest [] BL.empty + r <- readChan' resp + assertBool "Unexpected send response" (isNothing r) + + let headers = [("receipt","1"), + ("content-type","text/plain"), + ("content-length", show (BL.length testMsg)) + ] + send con dest headers testMsg + Just (Frame (SC RECEIPT) hs _) <- readChan' resp + let receipt = fromJust $ lookup "receipt-id" hs + assertEqual "Receipt for send command" receipt "1") + (disconnect con []) + +testSubscribe = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do send con dest [] testMsg + subscribe con dest subid [] + Just (Frame (SC cmd) hs body) <- readChan' resp + unsubscribe con subid [] + assertEqual "Wrong broker response" cmd MESSAGE + assertEqual "Invalid message received" body testMsg + let sid = fromJust $ lookup "subscription" hs + assertEqual "Invalid subscription id" sid subid) + (disconnect con []) + +testUnsubscribe = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [] + send con dest [] testMsg + (Just (Frame (SC cmd) _ _)) <- readChan' resp + assertEqual "Wrong broker response" cmd MESSAGE + unsubscribe con subid [] + r <- readChan' resp + assertBool "Unexpected unsubscribe response" (isNothing r) + send con dest [] testMsg + r1 <- readChan' resp + assertBool "Unexpected response after send" (isNothing r1)) + (disconnect con []) + +testAck = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [("ack", "client")] + send con dest [] testMsg + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message received" body testMsg) + (disconnect con []) + + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [("ack", "client")] + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message received" body testMsg + let msgid = fromJust $ lookup "message-id" hs + ack con subid msgid []) + (disconnect con []) + + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [("ack", "client")] + r <- readChan' resp + assertBool "Unexpected response after ack" (isNothing r)) + (disconnect con []) + +testNack = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (unless (maximum (versions con) == (1,0)) $ do + subscribe con dest subid [("ack", "client-individual")] + send con dest [] testMsg + (Just (Frame (SC cmd) hs body)) <- readChan' resp + let msgid = fromJust $ lookup "message-id" hs + nack con subid msgid [] + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message received" body testMsg + let msgid = fromJust $ lookup "message-id" hs + ack con subid msgid []) + (disconnect con []) + +testCommit = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [] + begin con tid [] + send con dest [("transaction",tid)] testMsg + r <- readChan' resp + assertBool "Unexpected response before commit" (isNothing r) + commit con tid [] + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message received" body testMsg) + (disconnect con []) + +testAbort = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [] + begin con tid [] + send con dest [("transaction",tid)] testMsg + r <- readChan' resp + assertBool "Unexpected response before abort" (isNothing r) + abort con tid[] + r' <- readChan' resp + assertBool "Unexpected response after abort" (isNothing r)) + (disconnect con []) + +testReceipt = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [("receipt","1")] + readChan' resp >>= \f -> assertReceipt f "1" + begin con tid [("transaction",tid),("receipt","2")] + readChan' resp >>= \f -> assertReceipt f "2" + abort con tid [("transaction",tid),("receipt","3")] + readChan' resp >>= \f -> assertReceipt f "3") + (do disconnect con [("receipt","4")] + readChan' resp >>= \f -> assertReceipt f "4") + +testSendbeat = do + initDest dest subid + (con, resp, _) <- mkStomp + if maximum (versions con) > (1,0) then do + threadDelay (sendTimeout con * 5000) + E.catch + (do disconnect con [] + assertBool "Send heartbeat failure" False) + (\(e::StompException) -> + assertBool "Invalid exception" (isConError e)) + (con, resp, _) <- mkStomp + startSendBeat con + threadDelay (sendTimeout con * 5000) + disconnect con [] + else disconnect con [] + +testRecvbeat = do + initDest dest subid + (con, resp, excp) <- mkStomp + if maximum (versions con) > (1,0) then do + startRecvBeat con + threadDelay (recvTimeout con * 5000) + e <- readChan' excp + assertBool "Receive heartbeat failure" (isBrokerError e) + else disconnect con [] + where + isBrokerError (Just (BrokerError _)) = True + isBrokerError _ = False + +testErrors = do + E.catch + (do + connect "invalid uri" [] [] + assertBool "Connect test failure" False) + (\(e::StompException) -> + case e of + InvalidUri _ -> return () + _ -> assertBool "Invalid exception" False) + + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do send con "/1queue" [] testMsg + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid error response" cmd ERROR) + (disconnect con []) + +testEncode = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [("receipt",":\n\\")] + readChan' resp >>= \f -> assertReceipt f ":\n\\") + (disconnect con []) + +testNull = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do + subscribe con dest subid [] + send con dest [] msg + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message size" (BL.length body) len + send con dest [] testMsg + (Just (Frame (SC cmd) hs body')) <- readChan' resp + assertEqual "Invalid message" body' testMsg) + (disconnect con []) + where + msg = BL.pack (replicate 1024 'x' ++ "\x00" ++ replicate 1024 'y') + len = BL.length msg + +testLength = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [] + send con dest [] msg + (Just (Frame (SC cmd) hs body)) <- readChan' resp + assertEqual "Invalid message received" body msg) + (disconnect con []) + where + msg = BL.pack (replicate 1024 'x' ++ replicate 1024 'y' ++ "\x00\x00zz") + +testBig = do + initDest dest subid + (con, resp, _) <- mkStomp + E.finally + (do subscribe con dest subid [] + send con dest [("content-length", show size)] msg + (Just (Frame (SC cmd) hs body)) <- readChan'' resp + assertEqual "Invalid message size" (BL.length body) (fromIntegral size)) + (disconnect con []) + where + size = 1024 * 1024 * 8 + msg = BL.pack (replicate size 'x') + +-- | Make default stomp connection with communication channels +mkStomp :: IO (Connection, Chan Frame, Chan StompException) +mkStomp = do + con <- connect uri [(1,0),(1,1)] [("heart-beat","5000,5000")] + respChan <- newChan + excpChan <- newChan + setExcpHandler con (excpFun con excpChan) + startConsumer con (respFun con respChan) + return (con, respChan, excpChan) + where + excpFun con = writeChan + respFun con = writeChan + +-- | Check receipt frame +assertReceipt :: Maybe Frame -> String -> IO () +assertReceipt (Just (Frame (SC cmd) hs _)) r = do + assertEqual "RECEIPT response expected" cmd RECEIPT + let r' = fromJust $ lookup "receipt-id" hs + assertEqual "Wrong receipt value" r' r + +-- | Check exception as a connection error +isConError :: StompException -> Bool +isConError (ConnectionError _) = True +isConError _ = False + +-- | Default test uri +uri :: String +uri = "stomp://guest:guest@127.0.0.1:61613" + +-- | Default test queue destination +dest :: String +dest = "/queue/q" + +-- | Default test subscription id +subid :: String +subid = "0" + +-- | Default test transaction id +tid :: String +tid = "0" + +-- | Test message body +testMsg :: BL.ByteString +testMsg = BL.pack "test message" + +-- | Read from channel with 2 sec timeout +readChan' :: Chan a -> IO (Maybe a) +readChan' ch = timeout 2000000 (readChan ch) + +-- | Read from channel with 15 sec timeout +readChan'' :: Chan a -> IO (Maybe a) +readChan'' ch = timeout 15000000 (readChan ch) + +-- | Initialize the destination before test +initDest :: Destination -> Subscription -> IO () +initDest dest subid = do + (con, resp, _) <- mkStomp + subscribe con dest subid [] + _ <- whileJust (readChan' resp) + disconnect con [] + +-- | Recursively collect values contained in the Just +whileJust :: Monad m => m (Maybe a) -> m [a] +whileJust p = do + x <- p + case x of + Nothing -> return mzero + Just x -> do + xs <- whileJust p + return $ return x `mplus` xs