diff --git a/dbus.cabal b/dbus.cabal --- a/dbus.cabal +++ b/dbus.cabal @@ -1,5 +1,5 @@ name: dbus -version: 0.10.1 +version: 0.10.2 license: GPL-3 license-file: license.txt author: John Millikin diff --git a/lib/DBus/Client.hs b/lib/DBus/Client.hs --- a/lib/DBus/Client.hs +++ b/lib/DBus/Client.hs @@ -116,6 +116,7 @@ -- * Advanced connection options , ClientOptions , clientSocketOptions + , clientThreadRunner , defaultClientOptions , connectWith ) where @@ -164,6 +165,16 @@ -- | Options for the underlying socket, for advanced use cases. See -- the "DBus.Socket" module. clientSocketOptions :: DBus.Socket.SocketOptions t + + -- | A function to run the client thread. The provided IO computation + -- should be called repeatedly; each time it is called, it will process + -- one incoming message. + -- + -- The provided computation will throw a 'ClientError' if it fails to + -- process an incoming message, or if the connection is lost. + -- + -- The default implementation is 'forever'. + , clientThreadRunner :: IO () -> IO () } type Callback = (ReceivedMessage -> IO ()) @@ -246,10 +257,12 @@ signalHandlers <- newIORef [] objects <- newIORef Data.Map.empty + let threadRunner = clientThreadRunner opts + clientMVar <- newEmptyMVar threadID <- forkIO $ do client <- readMVar clientMVar - mainLoop client + threadRunner (mainLoop client) let client = Client { clientSocket = sock @@ -273,6 +286,7 @@ defaultClientOptions :: ClientOptions SocketTransport defaultClientOptions = ClientOptions { clientSocketOptions = DBus.Socket.defaultSocketOptions + , clientThreadRunner = forever } -- | Stop a 'Client''s callback thread and close its underlying socket. @@ -293,7 +307,7 @@ DBus.Socket.close (clientSocket client) mainLoop :: Client -> IO () -mainLoop client = forever $ do +mainLoop client = do let sock = clientSocket client received <- Control.Exception.try (DBus.Socket.receive sock) diff --git a/lib/DBus/Transport.hs b/lib/DBus/Transport.hs --- a/lib/DBus/Transport.hs +++ b/lib/DBus/Transport.hs @@ -158,10 +158,16 @@ loop builder (n - Data.ByteString.length chunk) else do chunk <- recv s n - let builder = Builder.append acc (Builder.fromByteString chunk) - if Data.ByteString.length chunk == n - then return (Builder.toByteString builder) - else loop builder (n - Data.ByteString.length chunk) + case Data.ByteString.length chunk of + -- Unexpected end of connection; maybe the remote end went away. + -- Return what we've got so far. + 0 -> return (Builder.toByteString acc) + + len -> do + let builder = Builder.append acc (Builder.fromByteString chunk) + if len == n + then return (Builder.toByteString builder) + else loop builder (n - Data.ByteString.length chunk) instance TransportOpen SocketTransport where transportOpen _ a = case addressMethod a of diff --git a/lib/DBus/Wire.hs b/lib/DBus/Wire.hs --- a/lib/DBus/Wire.hs +++ b/lib/DBus/Wire.hs @@ -570,8 +570,11 @@ unmarshalMessageM :: Monad m => (Int -> m ByteString) -> m (Either UnmarshalError ReceivedMessage) unmarshalMessageM getBytes' = runErrorT $ do - let getBytes = ErrorT . liftM Right . getBytes' - + let getBytes count = do + bytes <- ErrorT (liftM Right (getBytes' count)) + if Data.ByteString.length bytes < count + then throwErrorT (UnmarshalError "Unexpected end of input while parsing message header.") + else return bytes let Just fixedSig = parseSignature "yyyyuuu" fixedBytes <- getBytes 16 @@ -671,9 +674,17 @@ require label _ = throwErrorM label unmarshalMessage :: ByteString -> Either UnmarshalError ReceivedMessage -unmarshalMessage bytes = case Get.runGet (unmarshalMessageM Get.getByteString) bytes of - Left err -> Left (UnmarshalError err) - Right x -> x +unmarshalMessage bytes = checkError (Get.runGet get bytes) where + get = unmarshalMessageM getBytes + + -- wrap getByteString, so it will behave like transportGet and return + -- a truncated result on EOF instead of throwing an exception. + getBytes count = do + remaining <- Get.remaining + Get.getByteString (min remaining count) + + checkError (Left err) = Left (UnmarshalError err) + checkError (Right x) = x untilM :: Monad m => m Bool -> m a -> m [a] untilM test comp = do diff --git a/tests/DBusTests.hs b/tests/DBusTests.hs --- a/tests/DBusTests.hs +++ b/tests/DBusTests.hs @@ -35,6 +35,7 @@ import DBusTests.Signature import DBusTests.Transport import DBusTests.Variant +import DBusTests.Wire -- import all dbus modules here to ensure they show up in the coverage report, -- even if not tested. @@ -64,6 +65,7 @@ , test_Socket , test_Transport , test_Variant + , test_Wire ] main :: IO () diff --git a/tests/DBusTests/Transport.hs b/tests/DBusTests/Transport.hs --- a/tests/DBusTests/Transport.hs +++ b/tests/DBusTests/Transport.hs @@ -42,6 +42,7 @@ test_TransportListen test_TransportAccept test_TransportSendReceive + test_HandleLostConnection test_TransportOpen :: Suite test_TransportOpen = suite "transportOpen" @@ -244,6 +245,22 @@ liftIO (transportPut t sentBytes) bytes <- liftIO (readMVar var) $assert (equal bytes sentBytes) + +test_HandleLostConnection :: Test +test_HandleLostConnection = assertions "handle-lost-connection" $ do + (addr, networkSocket) <- listenRandomIPv4 + afterTest (N.sClose networkSocket) + + _ <- liftIO $ forkIO $ do + (s, _) <- NS.accept networkSocket + sendAll s "123" + NS.sClose s + + t <- liftIO (transportOpen socketTransportOptions addr) + afterTest (transportClose t) + + bytes <- liftIO (transportGet t 4) + $assert (equal bytes "123") test_ListenUnknown :: Test test_ListenUnknown = assertions "unknown" $ do diff --git a/tests/DBusTests/Wire.hs b/tests/DBusTests/Wire.hs new file mode 100644 --- /dev/null +++ b/tests/DBusTests/Wire.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +-- Copyright (C) 2012 John Millikin +-- +-- This program is free software: you can redistribute it and/or modify +-- it under the terms of the GNU General Public License as published by +-- the Free Software Foundation, either version 3 of the License, or +-- any later version. +-- +-- This program is distributed in the hope that it will be useful, +-- but WITHOUT ANY WARRANTY; without even the implied warranty of +-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +-- GNU General Public License for more details. +-- +-- You should have received a copy of the GNU General Public License +-- along with this program. If not, see . + +module DBusTests.Wire (test_Wire) where + +import Test.Chell + +import qualified Data.ByteString.Char8 () + +import DBus + +test_Wire :: Suite +test_Wire = suite "Wire" + test_Unmarshal + +test_Unmarshal :: Suite +test_Unmarshal = suite "unmarshal" + test_UnmarshalUnexpectedEof + +test_UnmarshalUnexpectedEof :: Test +test_UnmarshalUnexpectedEof = assertions "unexpected-eof" $ do + let unmarshaled = unmarshal "0" + $assert (left unmarshaled) + + let Left err = unmarshaled + $assert (equal + (unmarshalErrorMessage err) + "Unexpected end of input while parsing message header.")