2 files changed, 20 insertions, 9 deletions

diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index 4db055ea..55aab3d4 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -646,10 +646,12 @@ handlePeer st src snk = do
     debugL $ "(P) connected " <++> name
     session <- newPeerSession session_factory sock
 
-    finally ( src $= parseBytes def $$ fromPeer sock session )
+    didClose <- newTVarIO False
+    finally ( src $= parseBytes def $$ fromPeer sock session didClose )
         $ do 
             debugL $ "(P) disconnected " <++> name
-            closePeerSession session
+            didc <- readTVarIO didClose
+            when (not didc) $ closePeerSession session
 
 
 handlePeerPresence session stanza False = do
@@ -912,8 +914,8 @@ handlePeerIQGet session tag = do
         req -> unhandledGet req
 
 fromPeer :: (MonadThrow m,MonadIO m, JabberPeerSession session) => 
-    RestrictedSocket -> session -> Sink XML.Event m ()
-fromPeer sock session = doNestingXML $ do
+    RestrictedSocket -> session -> TVar Bool -> Sink XML.Event m ()
+fromPeer sock session didClose = doNestingXML $ do
     let log  = liftIO . debugL . ("(P) " <++>)
     withXML $ \begindoc -> do
     when (begindoc==EventBeginDocument) $ do
@@ -922,7 +924,9 @@ fromPeer sock session = doNestingXML $ do
     withJust (elementAttrs "stream" xml) $ \stream_attrs -> do
     log $ "stream atributes: " <++> bshow stream_attrs
 
-    let doTimeout = Thunk (closePeerSession session)
+    let doTimeout = Thunk $ do 
+                        atomically $ writeTVar didClose True
+                        closePeerSession session
 
     fix $ \loop -> do
         log "waiting for stanza."
diff --git a/Presence/main.hs b/Presence/main.hs
index 45b8729a..f0dfc9bd 100644
--- a/Presence/main.hs
+++ b/Presence/main.hs
@@ -434,7 +434,11 @@ instance JabberPeerSession PeerSession where
         me <- fmap (RemotePeer . withoutPort) (getPeerName sock)
         debugL $ "PEER SESSION: open "<++>showPeer me
         let remotes = remoteUsers state
-        jids <- atomically $ getRefFromMap remotes me (newTVar MM.empty) return
+        (jids,us) <- atomically $ do
+                jids <- getRefFromMap remotes me (newTVar MM.empty) return
+                us <- readTVar (activeUsers state)
+                return (jids,map tupleToJID . Set.toList . Map.keysSet $ us)
+        forM_ us $ sendProbes state (Just me)
         return $ PeerSession jids me state
 
     closePeerSession session  = do
@@ -710,7 +714,7 @@ update_presence locals_greedy subscribers state getStatus =
         sendPresence chan jid status
       debugL $ bshow jid <++> " " <++> bshow status
 
-sendProbes state jid = do
+sendProbes state mbpeer jid = do
     debugL $ "sending probes for " <++> bshow jid
     withJust (name jid) $ \user -> do
     let parseHostNameJID' str = do
@@ -718,7 +722,10 @@ sendProbes state jid = do
                    (fmap Just . parseHostNameJID $ str)
     buddies <- do
         buddies <- ConfigFiles.getBuddies user
-        fmap catMaybes (mapM parseHostNameJID' buddies)
+        buddies' <- fmap catMaybes (mapM parseHostNameJID' buddies)
+        case mbpeer of
+          Nothing -> return buddies'
+          Just p -> return (filter (\jid-> peer jid == p) buddies')
     debugL $ "buddies for " <++> bshow jid <++> " = " <++> bshow buddies
     wanted <- do
         wanted <- ConfigFiles.getSolicited user
@@ -762,7 +769,7 @@ track_login host state e = do
     update_presence locals_greedy subs departures $ const Offline
     update_presence locals_greedy subs arrivals   $ matchResource active_users tty
     forM_ arrivals
-          $ sendProbes state
+          $ sendProbes state Nothing
 
 on_chvt state vtnum = do
     let tty = L.snoc "tty" $ intToDigit (fromIntegral vtnum)