1 files changed, 11 insertions, 6 deletions

diff --git a/Presence/SendMessage.hs b/Presence/SendMessage.hs
index 6e8ea2b9..d1db7a4f 100644
--- a/Presence/SendMessage.hs
+++ b/Presence/SendMessage.hs
@@ -1,8 +1,10 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleContexts #-}
 module SendMessage
     ( sendMessage
     , CommandCache(..)
+    , ThreadChannelCommand(..)
     , newOutgoingConnections
     , OutgoingConnections
     ) where
@@ -47,7 +49,7 @@ import Network.Socket
 import XMPPTypes (Peer(..),discardPort,peerAddr,withPort,socketFamily)
 import SocketLike
 import ServerC (packetSink)
-import ControlMaybe (handleIO,handleIO_)
+import ControlMaybe 
 import Data.Conduit (Sink,Source)
 import qualified Data.ByteString as S (ByteString)
 import XMLToByteStrings
@@ -64,11 +66,12 @@ modifyIORef' ref f = do
 
 class CommandCache cache where
     type CacheableCommand cache
-
     emptyCache :: cache
-
     updateCache :: CacheableCommand cache -> cache -> cache
 
+class ThreadChannelCommand cmd where
+    isQuitCommand :: cmd -> Bool
+
 
 data OutgoingConnections cache = 
         OutgoingConnections (TVar (Map Peer (TChan (CacheableCommand cache), ThreadId)))
@@ -91,7 +94,7 @@ newOutgoingConnections interpretCommands = do
 
 
 sendMessage
-  :: CommandCache a =>
+  :: (CommandCache a, ThreadChannelCommand (CacheableCommand a)) =>
      OutgoingConnections a -> CacheableCommand a -> Peer -> IO ()
 sendMessage (OutgoingConnections cons interpretCommands) msg peer0 = do
     let peer = discardPort peer0
@@ -136,14 +139,16 @@ connect_to_server chan peer toPeer = (>> return ()) . runMaybeT $ do
     fix $ \sendmsgs -> do
         connected <- liftIO . async $ connect' (peerAddr peer) port
 
-        sock <- MaybeT . fix $ \loop -> do
+        msock <- MaybeT . fix $ \loop -> do
             e <- atomically $ orElse
                     (fmap Right $ waitSTM connected)
                     (fmap Left  $ readTChan chan)
             case e of
+                Left cmd | isQuitCommand cmd -> return Nothing
                 Left cmd   -> cacheCmd cmd cached >> loop
-                Right sock -> return sock
+                Right sock -> return (Just sock)
 
+        withJust msock $ \sock -> do
         retry <- do
             (cache,snk) <- liftIO $ do
                 h <- socketToHandle sock ReadWriteMode