4 files changed, 19 insertions, 7 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
diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index 2c01d456..3caad740 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -885,6 +885,10 @@ instance CommandCache CachedMessages where
     updateCache (Rejection from to) cache =
                 cache { approvals= mmInsert (False,from) to $ approvals cache }
 
+instance ThreadChannelCommand OutBoundMessage where
+    isQuitCommand Disconnect = True
+    isQuitCommand _          = False
+
 mmInsert val key mm = Map.alter f key mm
  where 
     f Nothing    = Just $ Set.singleton val
@@ -993,7 +997,8 @@ toPeer sock cache chan fail = do
           Rejection from to -> do
             liftIO . debugL $ "sending rejection "<++>bshow (from,to)
             sendApproval False from to
-        loop
+          Disconnect -> return ()
+        when (not . isQuitCommand $ event) loop
     send goodbyePeer
 
 
diff --git a/Presence/XMPPTypes.hs b/Presence/XMPPTypes.hs
index 8d0bd242..6b0d9ed1 100644
--- a/Presence/XMPPTypes.hs
+++ b/Presence/XMPPTypes.hs
@@ -244,6 +244,7 @@ data OutBoundMessage = OutBoundPresence Presence
                         | Solicitation  JID JID
                         | Approval JID JID
                         | Rejection JID JID
+                        | Disconnect
  deriving Prelude.Show
 
 getNamesForPeer :: Peer -> IO [S.ByteString]
diff --git a/Presence/main.hs b/Presence/main.hs
index dedd546a..593bb4bb 100644
--- a/Presence/main.hs
+++ b/Presence/main.hs
@@ -385,6 +385,7 @@ instance JabberPeerSession PeerSession where
         let offline jid = Presence jid Offline
         unrefFromMap (remoteUsers . peer_global $ session) (peer_name session)
             $ do
+                sendPeerMessage session Disconnect
                 debugStr ("unrefFromMap!")
                 js <- fmap (MM.toAscList) (readTVarIO . announced $ session)
                 forM_ js $ \(u,rs) -> do