7 files changed, 85 insertions, 69 deletions

diff --git a/Connection/Tox.hs b/Connection/Tox.hs
index 0c2f281f..03ffbc25 100644
--- a/Connection/Tox.hs
+++ b/Connection/Tox.hs
@@ -35,10 +35,10 @@ import System.IO
 
 
 
-data Parameters = Parameters
+data Parameters extra = Parameters
     { -- | Various Tox transports and clients.
       dhtRouting :: Routing
-    , roster     :: ContactInfo
+    , roster     :: ContactInfo extra
     , sessions   :: NetCryptoSessions
     , dhtClient  :: DHT.Client
       -- | Thread to be forked when a connection is established.
@@ -127,7 +127,7 @@ lookupForPolicyChange conmap k policy = do
 callbackId :: Int
 callbackId = 1
 
-lookupContact :: Key -> ContactInfo -> STM (Maybe (SecretKey,Contact))
+lookupContact :: Key -> ContactInfo extra -> STM (Maybe (SecretKey,Contact))
 lookupContact (Key me them) ContactInfo{accounts} = do
     acnts <- readTVar accounts
     fmap join $ forM (HashMap.lookup me acnts) $ \Account{userSecret,contacts} -> do
@@ -136,7 +136,7 @@ lookupContact (Key me them) ContactInfo{accounts} = do
             return (userSecret,c)
 
 -- | This function will fork threads as necessary.
-setToxPolicy :: Parameters
+setToxPolicy :: Parameters extra
              -> TVar (Map.Map Key SessionState)
              -> Key
              -> Policy
@@ -249,7 +249,7 @@ stringToKey_ s = let (xs,ys) = break (==':') s
                                        them <- readMaybe (drop 1 ys)
                                        return $ Key me them
 
-toxManager :: Parameters -> IO (Manager ToxProgress Key)
+toxManager :: Parameters extra -> IO (Manager ToxProgress Key)
 toxManager params = do
     conmap <- newTVarIO Map.empty
     return Manager
diff --git a/ToxManager.hs b/ToxManager.hs
index 460b2fe5..cd835983 100644
--- a/ToxManager.hs
+++ b/ToxManager.hs
@@ -11,8 +11,8 @@ import Control.Concurrent.STM
 import Control.Monad
 import Crypto.Tox
 import qualified Data.HashMap.Strict         as HashMap
+import qualified Data.Map                    as Map
 import Data.Maybe
-import qualified Data.Set                    as Set
 import qualified Data.Text                   as T
 import Data.Time.Clock.POSIX
 import Network.Address
@@ -38,10 +38,11 @@ import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
 
-toxAnnounceSendData :: Tox.Tox -> PublicKey
-                               -> Nonce32
-                               -> Maybe Tox.NodeInfo
-                               -> IO (Maybe (Tox.Rendezvous, Tox.AnnounceResponse))
+toxAnnounceSendData :: Tox.Tox JabberClients
+                        -> PublicKey
+                        -> Nonce32
+                        -> Maybe Tox.NodeInfo
+                        -> IO (Maybe (Tox.Rendezvous, Tox.AnnounceResponse))
 toxAnnounceSendData tox pubkey token = \case
     Just ni -> Tox.putRendezvous (Tox.onionTimeout tox)
                                  (Tox.toxCryptoKeys tox)
@@ -56,7 +57,11 @@ toxAnnounceSendData tox pubkey token = \case
 --
 -- These hooks will be invoked in order to connect to *.tox hosts in a user's
 -- XMPP roster.
-toxman :: Announcer -> [(String,TVar (BucketList Tox.NodeInfo))] -> Tox.Tox -> PresenceState -> ToxManager ConnectionKey
+toxman :: Announcer
+            -> [(String,TVar (BucketList Tox.NodeInfo))]
+            -> Tox.Tox JabberClients
+            -> PresenceState
+            -> ToxManager ConnectionKey
 toxman announcer toxbkts tox presence = ToxManager
     { activateAccount        = \k pubname seckey -> do
         hPutStrLn stderr $ "toxman ACTIVATE " ++ show (Tox.key2id $ toPublic seckey)
@@ -65,11 +70,12 @@ toxman announcer toxbkts tox presence = ToxManager
             pubid = Tox.key2id pub
         (acnt,newlyActive) <- atomically $ do
             macnt <- HashMap.lookup pubid <$> readTVar accounts
-            acnt <- maybe (newAccount seckey) return macnt
-            rs <- readTVar $ clientRefs acnt
-            writeTVar (clientRefs acnt) $! Set.insert k rs
+            acnt <- maybe (newAccount seckey Map.empty) return macnt
+            rs <- readTVar $ accountExtra acnt
+            perclient <- initPerClient
+            writeTVar (accountExtra acnt) $! Map.insert k perclient rs
             modifyTVar accounts (HashMap.insert pubid acnt)
-            if not (Set.null rs)
+            if not (Map.null rs)
                 then return (acnt,Nothing)
                 else return (acnt,Just $ \nid -> foldr interweave []
                                                  . map (R.kclosest (searchSpace (toxQSearch tox))
@@ -101,11 +107,11 @@ toxman announcer toxbkts tox presence = ToxManager
         forM mpubid $ \pubid -> do
         refs <- do
             macnt <- HashMap.lookup pubid <$> readTVar accounts
-            rs <- fromMaybe Set.empty <$> mapM (readTVar . clientRefs) macnt
+            rs <- fromMaybe Map.empty <$> mapM (readTVar . accountExtra) macnt
             forM_ macnt $ \acnt -> do
-                modifyTVar' (clientRefs acnt) $ Set.delete k
+                modifyTVar' (accountExtra acnt) $ Map.delete k
             return rs
-        if (Set.null $ refs Set.\\ Set.singleton k) then do
+        if (Map.null $ Map.delete k refs) then do
             --   TODO
             --   If this is the last reference to a non-connected contact:
             --      Stop the recurring search for that contact
diff --git a/ToxToXMPP.hs b/ToxToXMPP.hs
index 9391a232..5495f5ad 100644
--- a/ToxToXMPP.hs
+++ b/ToxToXMPP.hs
@@ -100,11 +100,21 @@ key2jid nospam key = T.pack $ show $ NoSpamId nsp key
     nlo = fromIntegral (0x0FFFF .&. nospam) :: Word16
     nhi = fromIntegral (0x0FFFF .&. (nospam `shiftR` 16)) :: Word16
 
+type JabberClients = Map.Map ConnectionKey PerClient
+
+data PerClient = PerClient
+    {
+    }
+
+initPerClient :: STM PerClient
+initPerClient = do
+    return PerClient {}
+
 data ToxToXMPP = ToxToXMPP
     { txAnnouncer :: Announcer
-    , txAccount   :: Account
+    , txAccount   :: Account JabberClients
     , txPresence  :: PresenceState
-    , txTox       :: Tox
+    , txTox       :: Tox JabberClients
     }
 
 dispatch :: ToxToXMPP -> ContactEvent -> IO ()
@@ -119,8 +129,8 @@ dispatch tx (OnionRouted  theirkey (OnionFriendRequest fr) ) = do
                   , txAccount   = acnt
                   , txPresence  = st } = tx
     k2c <- atomically $ do
-        refs <- readTVar (clientRefs acnt)
-        k2c <- Map.filterWithKey (\k _ -> k `Set.member` refs) <$> readTVar (keyToChan st)
+        refs <- readTVar (accountExtra acnt)
+        k2c <- Map.filterWithKey (\k _ -> isJust $ k `Map.lookup` refs) <$> readTVar (keyToChan st)
         clients <- readTVar (clients st)
         return $ Map.intersectionWith (,) k2c clients
     -- TODO: Below we're using a hard coded default as their jabber user id.
@@ -204,7 +214,7 @@ stopConnecting ToxToXMPP{txAnnouncer=announcer,txAccount=acnt} them = do
     akey <- akeyDHTKeyShare announcer me them
     cancel announcer akey
 
-forkAccountWatcher :: Account -> Tox -> PresenceState -> Announcer -> IO ThreadId
+forkAccountWatcher :: Account JabberClients -> Tox JabberClients -> PresenceState -> Announcer -> IO ThreadId
 forkAccountWatcher acc tox st announcer = forkIO $ do
     myThreadId >>= flip labelThread ("tox-xmpp:"
                                       ++ show (key2id $ toPublic $ userSecret acc))
@@ -220,13 +230,13 @@ forkAccountWatcher acc tox st announcer = forkIO $ do
     forM_ (HashMap.toList cs) $ \(them,c) -> do
         startConnecting0 tx (id2key them) c
 
-    -- Loop endlessly until clientRefs is null.
+    -- Loop endlessly until accountExtra is null.
     fix $ \loop -> do
     mev <- atomically $
             (Just <$> readTChan chan)
             `orElse` do
-                refs <- readTVar $ clientRefs acc
-                check $ Set.null refs
+                refs <- readTVar $ accountExtra acc
+                check $ Map.null refs
                 return Nothing
 
     forM_ mev $ \ev -> dispatch tx ev >> loop
@@ -236,7 +246,7 @@ forkAccountWatcher acc tox st announcer = forkIO $ do
         stopConnecting tx (id2key them)
 
 
-toxQSearch :: Tox.Tox -> Search Tox.NodeId (IP, PortNumber) Nonce32 Tox.NodeInfo Tox.Rendezvous
+toxQSearch :: Tox extra -> Search Tox.NodeId (IP, PortNumber) Nonce32 Tox.NodeInfo Tox.Rendezvous
 toxQSearch tox = Tox.toxidSearch (Tox.onionTimeout tox) (Tox.toxCryptoKeys tox) (Tox.toxOnion tox)
 
 toxAnnounceInterval :: POSIXTime
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index d4216dae..d0d8678d 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -312,13 +312,13 @@ data Session = Session
     , swarms            :: Mainline.SwarmsDatabase
     , cryptosessions    :: Tox.NetCryptoSessions
     , toxkeys           :: TVar Tox.AnnouncedKeys
-    , roster            :: Tox.ContactInfo
+    , roster            :: Tox.ContactInfo JabberClients
     , announceToLan     :: IO ()
     , connectionManager :: Maybe ConnectionManager
     , onionRouter       :: OnionRouter
     , announcer         :: Announcer
     , signalQuit        :: IO ()
-    , mbTox             :: Maybe Tox.Tox
+    , mbTox             :: Maybe (Tox.Tox JabberClients)
     }
 
 exceptionsToClient :: ClientHandle -> IO () -> IO ()
@@ -555,7 +555,7 @@ clientSession s@Session{..} sock cnum h = do
                         let pubkey = toPublic secret
                         oldks <- atomically $ do
                             ks <- myKeyPairs roster
-                            Tox.addContactInfo roster secret
+                            Tox.addContactInfo roster secret Map.empty
                             return ks
                         let asString = show . Tox.key2id
                         hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) oldks
@@ -584,7 +584,7 @@ clientSession s@Session{..} sock cnum h = do
                             pairs = map (toPair . f) mbSecs
                         oldks <- atomically $ do
                             oldks <- myKeyPairs roster
-                            forM pairs $ \(sk,_) -> Tox.addContactInfo roster sk
+                            forM pairs $ \(sk,_) -> Tox.addContactInfo roster sk Map.empty
                             return oldks
                         hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) oldks
                                                   ++ map (mappend " *" . show . Tox.key2id .snd) pairs
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 20302343..a13a4f10 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -248,7 +248,7 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
             in client
     return $ either mkclient mkclient tblvar handlers
 
-data Tox = Tox
+data Tox extra  = Tox
     { toxDHT            :: DHT.Client
     , toxOnion          :: Onion.Client RouteId
     , toxToRoute        :: Transport String Onion.AnnouncedRendezvous (PublicKey,Onion.OnionData)
@@ -260,17 +260,17 @@ data Tox = Tox
     , toxTokens         :: TVar SessionTokens
     , toxAnnouncedKeys  :: TVar Onion.AnnouncedKeys
     , toxOnionRoutes    :: OnionRouter
-    , toxContactInfo    :: ContactInfo
+    , toxContactInfo    :: ContactInfo extra
     , toxAnnounceToLan  :: IO ()
-    , toxMgr :: Manager ToxProgress Key
+    , toxMgr            :: Manager ToxProgress Key
     }
 
 -- | initiate a netcrypto session, blocking
-netCrypto :: Tox -> SecretKey -> PublicKey{-UserKey -} -> IO [NetCryptoSession]
+netCrypto :: Tox extra -> SecretKey -> PublicKey{-UserKey -} -> IO [NetCryptoSession]
 netCrypto tox myseckey theirpubkey = netCryptoWithBackoff 1000000 tox myseckey theirpubkey
 
 -- | helper for 'netCrypto', initiate a netcrypto session, retry after specified millisecs
-netCryptoWithBackoff :: Int -> Tox -> SecretKey -> PublicKey -> IO [NetCryptoSession]
+netCryptoWithBackoff :: Int -> Tox extra -> SecretKey -> PublicKey -> IO [NetCryptoSession]
 netCryptoWithBackoff millisecs tox myseckey theirpubkey = do
     let mykeyAsId = key2id (toPublic myseckey)
     -- TODO: check status of connection here:
@@ -358,7 +358,7 @@ netCryptoWithBackoff millisecs tox myseckey theirpubkey = do
                             return []
 
 -- | Create a DHTPublicKey packet to send to a remote contact.
-getContactInfo :: Tox -> IO DHT.DHTPublicKey
+getContactInfo :: Tox extra -> IO DHT.DHTPublicKey
 getContactInfo Tox{toxCryptoKeys,toxRouting} = join $ atomically $ do
     r4 <- readTVar $ DHT.routing4 toxRouting
     r6 <- readTVar $ DHT.routing6 toxRouting
@@ -416,7 +416,7 @@ newTox :: TVar Onion.AnnouncedKeys   -- ^ Store of announced keys we are a rende
        -> SockAddr                   -- ^ Bind-address to listen on.
        -> Maybe NetCryptoSessions    -- ^ State of all one-on-one Tox links.
        -> Maybe SecretKey            -- ^ Optional DHT secret key to use.
-       -> IO Tox
+       -> IO (Tox extra)
 newTox keydb addr mbSessionsState suppliedDHTKey = do
     (udp,sock) <- {- addVerbosity <$> -} udpTransport' addr
     (crypto0,sessionsState0) <- case mbSessionsState of
@@ -506,21 +506,21 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
         , toxMgr            = mgr
         }
 
-onionTimeout :: Tox -> DHT.TransactionId -> Onion.OnionDestination RouteId -> STM (Onion.OnionDestination RouteId, Int)
+onionTimeout :: Tox extra -> DHT.TransactionId -> Onion.OnionDestination RouteId -> STM (Onion.OnionDestination RouteId, Int)
 onionTimeout Tox { toxOnionRoutes = or } (DHT.TransactionId n8 _) od = lookupTimeout or n8 od
 
 routing4nodeInfo :: DHT.Routing -> IO NodeInfo
 routing4nodeInfo (DHT.routing4 -> tv) = R.thisNode <$> readTVarIO tv
 
-dnssdAnnounce :: Tox -> IO ()
+dnssdAnnounce :: Tox extra -> IO ()
 dnssdAnnounce (toxRouting -> r) = do
     ni <- routing4nodeInfo r
     announceToxService (nodePort ni) (nodeId ni)
 
-dnssdDiscover :: Tox -> NodeInfo -> IO ()
+dnssdDiscover :: Tox extra -> NodeInfo -> IO ()
 dnssdDiscover (toxDHT -> client) ni = void $ DHT.ping client ni
 
-forkTox :: Tox -> IO (IO (), [NodeInfo] -> [NodeInfo] -> IO (), [NodeInfo] -> [NodeInfo] -> IO ())
+forkTox :: Tox extra -> IO (IO (), [NodeInfo] -> [NodeInfo] -> IO (), [NodeInfo] -> [NodeInfo] -> IO ())
 forkTox tox = do
     _ <- forkListener "toxHandshakes" (toxHandshakes tox)
     _ <- forkListener "toxToRoute" (toxToRoute tox)
diff --git a/src/Network/Tox/ContactInfo.hs b/src/Network/Tox/ContactInfo.hs
index 9f29d587..64ea861b 100644
--- a/src/Network/Tox/ContactInfo.hs
+++ b/src/Network/Tox/ContactInfo.hs
@@ -19,16 +19,16 @@ import Network.Tox.NodeId            (id2key)
 import Network.Tox.Onion.Transport   as Onion
 import System.IO
 
-newtype ContactInfo = ContactInfo
+newtype ContactInfo extra = ContactInfo
       -- | Map our toxid public key to an Account record.
-    { accounts :: TVar (HashMap NodeId{-my userkey-} Account)
+    { accounts :: TVar (HashMap NodeId{-my userkey-} (Account extra))
     }
 
-data Account = Account
-    { userSecret :: SecretKey                     -- local secret key
-    , contacts   :: TVar (HashMap NodeId{-friend's userkey-} Contact) -- received contact info
-    , clientRefs :: TVar (Set ConnectionKey) -- reference count so we know when to send offline etc
-    , eventChan  :: TChan ContactEvent
+data Account extra = Account
+    { userSecret   :: SecretKey                     -- local secret key
+    , contacts     :: TVar (HashMap NodeId{-friend's userkey-} Contact) -- received contact info
+    , accountExtra :: TVar extra
+    , eventChan    :: TChan ContactEvent
     }
 
 data ContactEvent = OnionRouted  { contact :: PublicKey, onionRouted :: OnionData }
@@ -44,16 +44,16 @@ data Contact = Contact
     , contactPolicy        :: TVar (Maybe Connection.Policy)
     }
 
-newContactInfo :: IO ContactInfo
+newContactInfo :: IO (ContactInfo extra)
 newContactInfo = atomically $ ContactInfo <$> newTVar HashMap.empty
 
-myKeyPairs :: ContactInfo -> STM [(SecretKey,PublicKey)]
+myKeyPairs :: ContactInfo extra -> STM [(SecretKey,PublicKey)]
 myKeyPairs (ContactInfo accounts) = do
     acnts <- readTVar accounts
     forM (HashMap.toList acnts) $ \(nid, Account{userSecret}) -> do
         return (userSecret,id2key nid)
 
-updateContactInfo :: ContactInfo -> Onion.AnnouncedRendezvous -> (PublicKey,Onion.OnionData) -> IO ()
+updateContactInfo :: ContactInfo extra -> Onion.AnnouncedRendezvous -> (PublicKey,Onion.OnionData) -> IO ()
 updateContactInfo roster Onion.AnnouncedRendezvous{remoteUserKey} (localUserKey,omsg) = do
     hPutStrLn stderr "updateContactInfo!!!"
     now <- getPOSIXTime
@@ -69,7 +69,7 @@ initContact = Contact <$> newTVar Nothing
                       <*> newTVar Nothing
                       <*> newTVar Nothing
 
-updateAccount' :: PublicKey -> Account -> (Contact -> STM ()) -> STM ()
+updateAccount' :: PublicKey -> Account extra -> (Contact -> STM ()) -> STM ()
 updateAccount' remoteUserKey acc updater = do
     let rkey = key2id remoteUserKey
     cmap <- readTVar (contacts acc)
@@ -80,7 +80,7 @@ updateAccount' remoteUserKey acc updater = do
                            return contact
     updater contact
 
-updateAccount :: POSIXTime -> PublicKey -> Onion.OnionData -> Account -> STM ()
+updateAccount :: POSIXTime -> PublicKey -> Onion.OnionData -> Account extra -> STM ()
 updateAccount now remoteUserKey omsg acc = do
     updateAccount' remoteUserKey acc $ onionUpdate now omsg
     writeTChan (eventChan acc) $ OnionRouted remoteUserKey omsg
@@ -97,39 +97,39 @@ policyUpdate policy contact = writeTVar (contactPolicy contact) $ Just policy
 addrUpdate :: POSIXTime -> SockAddr -> Contact -> STM ()
 addrUpdate now addr contact = writeTVar (contactLastSeenAddr contact) $ Just (now,addr)
 
-setContactPolicy :: PublicKey -> Policy -> Account -> STM ()
+setContactPolicy :: PublicKey -> Policy -> Account extra -> STM ()
 setContactPolicy remoteUserKey policy acc = do
     updateAccount' remoteUserKey acc $ policyUpdate policy
     writeTChan (eventChan acc) $ PolicyChange remoteUserKey policy
 
-setContactAddr :: POSIXTime -> PublicKey -> SockAddr -> Account -> STM ()
+setContactAddr :: POSIXTime -> PublicKey -> SockAddr -> Account extra -> STM ()
 setContactAddr now remoteUserKey addr acc = do
     updateAccount' remoteUserKey acc $ addrUpdate now addr
     writeTChan (eventChan acc) $ AddrChange remoteUserKey addr
 
-setEstablished :: POSIXTime -> PublicKey -> Account -> STM ()
+setEstablished :: POSIXTime -> PublicKey -> Account extra -> STM ()
 setEstablished now remoteUserKey acc =
     writeTChan (eventChan acc) $ SessionEstablished remoteUserKey
 
-setTerminated :: POSIXTime -> PublicKey -> Account -> STM ()
+setTerminated :: POSIXTime -> PublicKey -> Account extra -> STM ()
 setTerminated now remoteUserKey acc =
     writeTChan (eventChan acc) $ SessionTerminated remoteUserKey
 
 
-addContactInfo :: ContactInfo -> SecretKey -> STM ()
-addContactInfo (ContactInfo as) sk = do
-    a <- newAccount sk
+addContactInfo :: ContactInfo extra -> SecretKey -> extra -> STM ()
+addContactInfo (ContactInfo as) sk extra = do
+    a <- newAccount sk extra
     modifyTVar' as $ HashMap.insert (key2id $ toPublic sk) a
 
-delContactInfo :: ContactInfo -> PublicKey -> STM ()
+delContactInfo :: ContactInfo extra -> PublicKey -> STM ()
 delContactInfo (ContactInfo as) pk = modifyTVar' as $ HashMap.delete (key2id pk)
 
-newAccount :: SecretKey -> STM Account
-newAccount sk = Account sk <$> newTVar HashMap.empty
-                           <*> newTVar Set.empty
-                           <*> newBroadcastTChan
+newAccount :: SecretKey -> extra -> STM (Account extra)
+newAccount sk extra = Account sk <$> newTVar HashMap.empty
+                                 <*> newTVar extra
+                                 <*> newBroadcastTChan
 
-dnsPresentation :: ContactInfo -> STM String
+dnsPresentation :: ContactInfo extra -> STM String
 dnsPresentation (ContactInfo accsvar) = do
     accs <- readTVar accsvar
     ms <- forM accs $ \Account { userSecret = sec, contacts = cvar } -> do
@@ -150,7 +150,7 @@ dnsPresentation1 (nid,dk) = unlines
 type LocalKey = NodeId
 type RemoteKey = NodeId
 
-friendRequests :: ContactInfo -> STM (HashMap LocalKey [(RemoteKey,DHT.FriendRequest)])
+friendRequests :: ContactInfo extra -> STM (HashMap LocalKey [(RemoteKey,DHT.FriendRequest)])
 friendRequests (ContactInfo roster) = do
     accs <- readTVar roster
     forM accs $ \Account { userSecret = sec, contacts = cvar } -> do
diff --git a/src/Network/Tox/DHT/Transport.hs b/src/Network/Tox/DHT/Transport.hs
index 5fdcd252..9ff1839c 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -262,7 +262,7 @@ data FriendRequest = FriendRequest
     { friendNoSpam      :: Word32
     , friendRequestText :: ByteString -- UTF8
     }
- deriving (Eq, Show)
+ deriving (Eq, Ord, Show)
 
 
 -- When sent as a DHT request packet (this is the data sent in the DHT request