3 files changed, 59 insertions, 33 deletions

diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index 9dfee14e..3d20a9b8 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE ViewPatterns        #-}
 module XMPP 
     ( module XMPPTypes
+    , module SocketLike
     , listenForXmppClients
     , listenForRemotePeers
     , seekRemotePeers
@@ -675,8 +676,6 @@ connect' addr port = do
     -}
     let getport (SockAddrInet port _) = port
         getport (SockAddrInet6 port _ _ _) = port
-    let withPort (SockAddrInet _ a)      port = SockAddrInet (toEnum port) a
-        withPort (SockAddrInet6 _ a b c) port = SockAddrInet6 (toEnum port) a b c
     let doException (SomeException e) = do
             L.putStrLn $  "\nFailed to reach "<++> showPeer (RemotePeer addr) <++> " on port "<++>bshow port<++>": " <++> bshow e
             return Nothing
diff --git a/Presence/XMPPTypes.hs b/Presence/XMPPTypes.hs
index 8af1018c..b46b4294 100644
--- a/Presence/XMPPTypes.hs
+++ b/Presence/XMPPTypes.hs
@@ -164,4 +164,8 @@ socketFamily (SockAddrInet _ _)      = AF_INET
 socketFamily (SockAddrInet6 _ _ _ _) = AF_INET6
 socketFamily (SockAddrUnix _)        = AF_UNIX
 
+withPort (SockAddrInet _ a)      port = SockAddrInet (toEnum port) a
+withPort (SockAddrInet6 _ a b c) port = SockAddrInet6 (toEnum port) a b c
+
+withoutPort = (`withPort` 0)
 
diff --git a/Presence/main.hs b/Presence/main.hs
index f17b8340..3182431f 100644
--- a/Presence/main.hs
+++ b/Presence/main.hs
@@ -56,6 +56,18 @@ import System.Environment
 import qualified Text.Show.ByteString as L
 import Network.Socket (Family(AF_INET,AF_INET6))
 
+type RefCount = Int
+
+data PresenceState = PresenceState 
+    { hostname :: Peer -- ByteString, TODO: remove this, its always LocalHost now
+    , currentTTY    :: TVar ByteString
+    , activeUsers   :: TVar (Set JID)
+    , subscriberMap :: TVar (Map JID (RefCount,TChan Presence)) -- UNUSED as yet
+    , localSubscriber :: TMVar (RefCount,TChan Presence) -- TODO: rename this, its not just locals
+                                                         -- ... or make a seperate channel for remotes
+    , remoteUsers :: TVar (Map Peer (RefCount,TVar (Set JID)))
+    }
+
 
 data ClientSession = ClientSession {
     localhost :: Peer, -- ByteString,
@@ -110,6 +122,7 @@ instance XMPPSession ClientSession where
 
 data PeerSession = PeerSession {
     announced :: TVar (Set JID),
+    peer_name :: Peer,
     peer_global :: PresenceState
 }
 instance XMPPSession PeerSession where
@@ -120,15 +133,19 @@ instance XMPPSession PeerSession where
     subscribe _ _   = error "subscribe on peer session?"
 
     newSession (PeerSessions state) sock = do
-        L.putStrLn $ "PEER SESSION: open"
-        jids <- newTVarIO Set.empty
-        return $ PeerSession jids state
+        me <- fmap (RemotePeer . withoutPort) (getPeerName sock)
+        L.putStrLn $ "PEER SESSION: open "<++>showPeer me
+        let remotes = remoteUsers state
+        jids <- atomically $ getRefFromMap remotes me (newTVar Set.empty) return
+        return $ PeerSession jids me state
 
     closeSession session  = do
-        L.putStrLn "PEER SESSION: close"
-        js <- fmap Set.toList (readTVarIO . announced $ session)
+        L.putStrLn $ "PEER SESSION: close "<++>showPeer (peer_name session)
         let offline jid = Presence jid Offline
-        forM_ js $ announcePresence session . offline
+        unrefFromMap (remoteUsers . peer_global $ session) (peer_name session)
+            $ do
+                js <- fmap Set.toList (readTVarIO . announced $ session)
+                forM_ js $ announcePresence session . offline
 
     announcePresence session (Presence jid status) = do
         (greedy,subs) <- atomically $ do
@@ -153,19 +170,35 @@ subscribeToChan tmvar =
                 putTMVar tmvar (1,chan)
                 return chan )
 
-subscribeToMap tvar jid = do
-            subs <- readTVar tvar
-            let mbchan       = Map.lookup jid subs
-            (chan,subs') <- 
-                do case mbchan of
-                    Nothing -> do
-                      newchan <- newTChan
-                      return (newchan, Map.insert jid (1,newchan) subs)
-                    Just (cnt,chan) -> do
-                      chan' <- dupTChan chan
-                      return (chan', Map.insert jid (cnt+1,chan) subs)
-            writeTVar tvar subs'
-            return chan
+getRefFromMap tvar key newObject copyObject = do
+    subs <- readTVar tvar
+    let mbobject       = Map.lookup key subs
+    (object,subs') <- 
+        do case mbobject of
+            Nothing -> do
+              newobject <- newObject
+              return (newobject, Map.insert key (1,newobject) subs)
+            Just (cnt,object) -> do
+              object' <- copyObject object
+              return (object', Map.insert key (cnt+1,object) subs)
+    writeTVar tvar subs'
+    return object
+
+unrefFromMap tvar key finalizer = do
+    vanished <- atomically $ do
+        omap <- readTVar tvar
+        let (r,omap') = Map.updateLookupWithKey unref key omap
+        writeTVar tvar omap'
+        return (isNothing r)
+    when vanished finalizer
+ where
+    unref key (cnt,object) =
+        if cnt==1 then Nothing
+                  else Just (cnt-1,object)
+    
+
+subscribeToMap tvar jid =
+    getRefFromMap tvar jid newTChan dupTChan
 
 matchResource tty jid = maybe Away (avail . (==tty)) $ resource jid
  where
@@ -188,23 +221,13 @@ update_presence locals_greedy subscribers state getStatus =
         sendPresence chan jid status
       putStrLn $ bshow jid <++> " " <++> bshow status
 
-type RefCount = Int
-
-data PresenceState = PresenceState 
-    { hostname :: Peer -- ByteString, TODO: remove this, its always LocalHost now
-    , currentTTY    :: TVar ByteString
-    , activeUsers   :: TVar (Set JID)
-    , subscriberMap :: TVar (Map JID (RefCount,TChan Presence)) -- UNUSED as yet
-    , localSubscriber :: TMVar (RefCount,TChan Presence) -- TODO: rename this, its not just locals
-                                                         -- ... or make a separte channel for remotes
-    }
-
 newPresenceState hostname = atomically $ do
     tty  <- newTVar ""
     us   <- newTVar (Set.empty)
     subs <- newTVar (Map.empty)
     locals_greedy <- newEmptyTMVar
-    return $ PresenceState hostname tty us subs locals_greedy
+    remotes <- newTVar (Map.empty)
+    return $ PresenceState hostname tty us subs locals_greedy remotes
 
 track_login host state e = do
 #ifndef NOUTMP