1 files changed, 374 insertions, 0 deletions

diff --git a/dht/src/Network/Tox/AggregateSession.hs b/dht/src/Network/Tox/AggregateSession.hs
new file mode 100644
index 00000000..8c728660
--- /dev/null
+++ b/dht/src/Network/Tox/AggregateSession.hs
@@ -0,0 +1,374 @@
+-- | This module aggregates all sessions to the same remote Tox contact into a
+-- single online/offline presence.  This allows multiple lossless links to the
+-- same identity at different addresses, or even to the same address.
+{-# LANGUAGE CPP             #-}
+{-# LANGUAGE GADTs           #-}
+{-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE PatternSynonyms #-}
+module Network.Tox.AggregateSession
+    ( AggregateSession
+    , newAggregateSession
+    , aggregateStatus
+    , checkCompatible
+    , compatibleKeys
+    , AddResult(..)
+    , addSession
+    , DelResult(..)
+    , delSession
+    , closeAll
+    , awaitAny
+    , dispatchMessage
+    ) where
+
+
+import Control.Concurrent.STM
+import Control.Concurrent.STM.TMChan
+import Control.Monad
+import Data.Dependent.Sum
+import Data.Function
+import qualified Data.IntMap.Strict  as IntMap
+         ;import Data.IntMap.Strict  (IntMap)
+import Data.List
+import Data.Time.Clock.POSIX
+import System.IO.Error
+
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent.Lifted
+import GHC.Conc                  (labelThread)
+#endif
+
+import Connection                   (Status (..))
+import Crypto.Tox                   (PublicKey, toPublic)
+import Data.Tox.Msg
+import Data.Wrapper.PSQInt          as PSQ
+import DPut
+import DebugTag
+import Network.QueryResponse
+import Network.Tox.Crypto.Transport
+import Network.Tox.DHT.Transport    (key2id)
+import Network.Tox.NodeId           (ToxProgress (..))
+import Network.Tox.Session
+
+-- | For each component session, we track the current status.
+data SingleCon = SingleCon
+    { singleSession :: Session                   -- ^ A component session.
+    , singleStatus  :: TVar (Status ToxProgress) -- ^ Either 'AwaitingSessionPacket' or 'Established'.
+    }
+
+-- | A collection of sessions between the same local and remote identities.
+data AggregateSession = AggregateSession
+    { -- | The set of component sessions indexed by their ID.
+      contactSession     :: TVar (IntMap SingleCon)
+      -- | Each inbound packets is written to this channel with the session ID
+      -- from which it came originally.
+    , contactChannel     :: TMChan (Int,CryptoMessage)
+      -- | The set of 'Established' sessions IDs.
+    , contactEstablished :: TVar (IntMap ())
+      -- | Callback for state-change notifications.
+    , notifyState        :: AggregateSession -> Session -> Status ToxProgress -> STM ()
+    }
+
+
+-- | Create a new empty aggregate session.  The argument is a callback to
+-- receive notifications when the new session changes status.  There are three
+-- possible status values:
+--
+--   [ Dormant ] - No pending or established sessions.
+--
+--   [ InProgress AwaitingSessionPacket ] - Sessions are pending, but none are
+--                                          fully established.
+--
+--   [ Established ] - At least one session is fully established and we can
+--                     send and receive packets via this aggregate.
+--
+-- The 'Session' object is provided to the callback so that it can determine the
+-- current remote and local identities for this AggregateSession.  It may not even
+-- be Established, so do not use it to send or receive packets.
+newAggregateSession :: (AggregateSession -> Session -> Status ToxProgress -> STM ())
+                        -> STM AggregateSession
+newAggregateSession notify = do
+    vimap <- newTVar IntMap.empty
+    chan <- newTMChan
+    vemap <- newTVar IntMap.empty
+    return AggregateSession
+        { contactSession     = vimap
+        , contactChannel     = chan
+        , contactEstablished = vemap
+        , notifyState        = notify
+        }
+
+-- | Information returned from 'addSession'.  Note that a value other than
+-- 'RejectedSession' does not mean there is any 'Established' session in the
+-- Aggregate.  Sessions are in 'AwaitingSessionPacket' state until a single
+-- packet is received from the remote end.
+data AddResult = FirstSession    -- ^ Initial connection with this contact.
+               | AddedSession    -- ^ Added another connection to active session.
+               | RejectedSession -- ^ Failed to add session (wrong contact / closed session).
+
+-- | The 'keepAlive' thread juggles three scheduled tasks.
+data KeepAliveEvents = DoTimeout        -- ^ A session timed-out, close it.
+                     | DoAlive          -- ^ Send a the keep-alive becon for a session.
+                     | DoRequestMissing -- ^ Detect and request lost packets.
+                     deriving Enum
+
+-- | This call loops until the provided sesison is closed or times out.  It
+-- monitors the provided (non-empty) priority queue for scheduled tasks (see
+-- 'KeepAliveEvents') to perform for the connection.
+keepAlive :: Session -> TVar (PSQ POSIXTime) -> IO ()
+keepAlive s q = do
+    myThreadId >>= flip labelThread
+        (intercalate "." ["beacon"
+                         , take 8 $ show $ key2id $ sTheirUserKey s
+                         , show $ sSessionID s])
+
+    let -- outPrint e   = dput XNetCrypto  $ shows (sSessionID s,sTheirAddr s) $ " <-- " ++ e
+        unexpected e = dput XUnexpected $ shows (sSessionID s,sTheirAddr s) $ " <-- " ++ e
+
+        doAlive = do
+            -- outPrint $ "Beacon"
+            sendMessage (sTransport s) () (Pkt ALIVE ==> ())
+
+        doRequestMissing = do
+            (ns,nmin) <- sMissingInbound s
+            -- outPrint $ "PacketRequest " ++ show (nmin,ns)
+            sendMessage (sTransport s) () (Pkt PacketRequest ==> MissingPackets ns)
+                `catchIOError` \e -> do
+                    unexpected $ "PacketRequest " ++ take 200 (show (nmin,length ns,ns))
+                    unexpected $ "PacketRequest: " ++ show e
+                    -- Quit thread by scheduling a timeout event.
+                    now <- getPOSIXTime
+                    atomically $ modifyTVar' q $ PSQ.insert (fromEnum DoTimeout) now
+
+        re tm again e io = do
+            io
+            atomically $ modifyTVar' q $ PSQ.insert (fromEnum e) tm
+            again
+
+        doEvent again now e = case e of
+            DoTimeout        -> do dput XNetCrypto $ "TIMEOUT: " ++ show (sSessionID s)
+                                   sClose s
+            DoAlive          -> re (now + 10) again e doAlive
+            DoRequestMissing -> re (now + 5)  again e doRequestMissing -- tox-core does this at 1 second intervals
+
+    fix $ \again -> do
+
+    now <- getPOSIXTime
+    join $ atomically $ do
+        PSQ.findMin <$> readTVar q >>= \case
+          Nothing           -> error "keepAlive: unexpected empty PSQ."
+          Just ( k :-> tm ) ->
+            return $ if now < tm then threadDelay (toMicroseconds $ tm - now) >> again
+                                 else doEvent again now (toEnum k)
+
+
+-- | This function forks two threads: the 'keepAlive' beacon-sending thread and
+-- a thread to read all packets from the provided 'Session' and forward them to
+-- 'contactChannel' for a containing 'AggregateSession'
+forkSession :: AggregateSession -> Session -> (Status ToxProgress -> STM ()) -> IO ThreadId
+forkSession c s setStatus = forkIO $ do
+    myThreadId >>= flip labelThread
+        (intercalate "." ["s"
+                         , take 8 $ show $ key2id $ sTheirUserKey s
+                         , show $ sSessionID s])
+
+    q <- atomically $ newTVar $ fromList
+            [ fromEnum DoAlive :-> 0
+            , fromEnum DoRequestMissing :-> 0
+            ]
+
+    let sendPacket :: CryptoMessage -> STM ()
+        sendPacket msg = writeTMChan (contactChannel c) (sSessionID s, msg)
+
+        inPrint e = dput XNetCrypto $ shows (sSessionID s,sTheirAddr s) $ " --> " ++ e
+
+        bump = do
+            -- inPrint $ "BUMP: " ++ show (sSessionID s)
+            now <- getPOSIXTime
+            atomically $ modifyTVar' q $ PSQ.insert (fromEnum DoTimeout) (now + 15)
+
+        onPacket body loop Nothing          = return ()
+        onPacket body loop (Just (Left e))  = inPrint e >> loop
+        onPacket body loop (Just (Right x)) = body loop x
+
+        awaitPacket body = fix $ awaitMessage (sTransport s) . onPacket body
+
+    atomically $ setStatus $ InProgress AwaitingSessionPacket
+    awaitPacket $ \_ (online,()) -> do
+        when (msgID online /= M ONLINE) $ do
+            inPrint $ "Unexpected initial packet: " ++ show (msgID online)
+        atomically $ do setStatus Established
+                        sendPacket online
+        bump
+        beacon <- forkIO $ keepAlive s q
+        awaitPacket $ \awaitNext (x,()) -> do
+            bump
+            case msgID x of
+                M ALIVE      -> return ()
+                M KillPacket -> sClose s
+                _            -> atomically $ sendPacket x
+            awaitNext
+        atomically $ setStatus Dormant
+        killThread beacon
+
+-- | Add a new session (in 'AwaitingSessionPacket' state) to the
+-- 'AggregateSession'.  If the supplied session is not compatible because it is
+-- between the wrong ToxIDs or because the AggregateSession is closed,
+-- 'RejectedSession' will be returned.  Otherwise, the operation is successful.
+--
+-- The status-change callback may be triggered by this call as the aggregate
+-- may transition from 'Dormant' (empty) to 'AwaitingSessionPacket' (at least
+-- one active session).
+addSession :: AggregateSession -> Session -> IO AddResult
+addSession c s = do
+    (result,mcon,replaced) <- atomically $ do
+        let them = sTheirUserKey s
+            me   = toPublic $ sOurKey s
+        compat <- checkCompatible me them c
+        let result = case compat of
+                        Nothing      -> FirstSession
+                        Just True    -> AddedSession
+                        Just False   -> RejectedSession
+        case result of
+            RejectedSession -> return (result,Nothing,Nothing)
+            _               -> do
+                                statvar <- newTVar Dormant
+                                imap <- readTVar (contactSession c)
+                                let con = SingleCon s statvar
+                                    s0 = IntMap.lookup (sSessionID s) imap
+                                    imap' = IntMap.insert (sSessionID s) con imap
+                                writeTVar (contactSession c) imap'
+                                return (result,Just con,s0)
+
+    mapM_ (sClose . singleSession) replaced
+    forM_ mcon $ \con ->
+        forkSession c s $ \progress -> do
+            writeTVar (singleStatus con) progress
+            emap <- readTVar (contactEstablished c)
+            emap' <- case progress of
+                Established -> do
+                    when (IntMap.null emap) $ notifyState c c s Established
+                    return $ IntMap.insert (sSessionID s) () emap
+                _ -> do
+                    let emap' = IntMap.delete (sSessionID s) emap
+                    when (IntMap.null emap' && not (IntMap.null emap)) $ do
+                        imap <- readTVar (contactSession c)
+                        notifyState c c s
+                            $ if IntMap.null imap then Dormant
+                                                  else InProgress AwaitingSessionPacket
+                    return emap'
+            writeTVar (contactEstablished c) emap'
+    return result
+
+-- | Information returned from 'delSession'.
+data DelResult = NoSession      -- ^ Contact is completely disconnected.
+               | DeletedSession -- ^ Connection removed but session remains active.
+
+-- | Close and remove the componenent session corresponding to the provided
+-- Session ID.
+--
+-- The status-change callback may be triggered as the aggregate may may
+-- transition to 'Dormant' (empty) or 'AwaitingSessionPacket' (if the last
+-- 'Established' session is closed).
+delSession :: AggregateSession -> Int -> IO DelResult
+delSession c sid = do
+    (con, r) <- atomically $ do
+        imap <- readTVar (contactSession c)
+        emap <- readTVar (contactEstablished c)
+        let emap' = IntMap.delete sid emap
+            imap' = IntMap.delete sid imap
+        case IntMap.toList emap of
+            (sid0,_):_ | IntMap.null emap'
+                       , let s = singleSession $ imap IntMap.! sid0
+                -> notifyState c c s
+                        $ if IntMap.null imap' then Dormant
+                                               else InProgress AwaitingSessionPacket
+            _   -> return ()
+        writeTVar (contactSession c) imap'
+        writeTVar (contactEstablished c) emap'
+        return ( IntMap.lookup sid imap, IntMap.null imap')
+    mapM_ (sClose . singleSession) con
+    return $ if r then NoSession
+                  else DeletedSession
+
+-- | Send a packet to one or all of the component sessions in the aggregate.
+dispatchMessage :: AggregateSession -> Maybe Int -- ^ 'Nothing' to broadcast, otherwise SessionID.
+                                    -> CryptoMessage -> IO ()
+dispatchMessage c msid msg = join $ atomically $ do
+    imap <- readTVar (contactSession c)
+    let go = case msid of Nothing  -> forM_ imap
+                          Just sid -> forM_ (IntMap.lookup sid imap)
+    return $ go $ \con -> sendMessage (sTransport $ singleSession con) () msg
+
+-- | Retry until:
+--
+--      * a packet arrives (with component session ID) arrives.
+--
+--      * the 'AggregateSession' is closed with 'closeAll'.
+awaitAny :: AggregateSession -> STM (Maybe (Int,CryptoMessage))
+awaitAny c = readTMChan (contactChannel c)
+
+-- | Close all connections associated with the aggregate.  No new sessions will
+-- be accepted after this, and the notify callback will be informed that we've
+-- transitioned to 'Dormant'.
+closeAll :: AggregateSession -> IO ()
+closeAll c = join $ atomically $ do
+    imap <- readTVar (contactSession c)
+    closeTMChan (contactChannel c)
+    return $ forM_ (IntMap.keys imap) $ \sid -> delSession c sid
+
+-- | Query the current status of the aggregate, there are three possible
+-- values:
+--
+--   [ Dormant ] - No pending or established sessions.
+--
+--   [ InProgress AwaitingSessionPacket ] - Sessions are pending, but none are
+--                                          fully established.
+--
+--   [ Established ] - At least one session is fully established and we can
+--                     send and receive packets via this aggregate.
+--
+aggregateStatus :: AggregateSession -> STM (Status ToxProgress)
+aggregateStatus c = do
+    isclosed <- isClosedTMChan (contactChannel c)
+    imap <- readTVar (contactSession c)
+    emap <- readTVar (contactEstablished c)
+    return $ case () of
+        _ | isclosed               -> Dormant
+          | not (IntMap.null emap) -> Established
+          | not (IntMap.null imap) -> InProgress AwaitingSessionPacket
+          | otherwise              -> Dormant
+
+-- | Query whether the supplied ToxID keys are compatible with this aggregate.
+--
+--   [ Nothing ]    Any keys would be compatible because there is not yet any
+--                  sessions in progress.
+--
+--   [ Just True ]  The supplied keys match the session in progress.
+--
+--   [ Just False ] The supplied keys are incompatible.
+checkCompatible :: PublicKey    -- ^ Local Tox key (for which we know the secret).
+                   -> PublicKey -- ^ Remote Tox key.
+                   -> AggregateSession -> STM (Maybe Bool)
+checkCompatible me them c = do
+    isclosed <- isClosedTMChan (contactChannel c)
+    imap <- readTVar (contactSession c)
+    return $ case IntMap.elems imap of
+                    _ | isclosed -> Just False -- All keys are incompatible (closed).
+                    con:_        -> Just $ sTheirUserKey (singleSession con)         == them
+                                           && toPublic (sOurKey $ singleSession con) == me
+                    []           -> Nothing
+
+-- | Returns the local and remote keys that are compatible with this aggregate.
+-- If 'Nothing' Is returned, then either no key is compatible ('closeAll' was
+-- called) or all keys are compatible because no sessions have been associated.
+compatibleKeys :: AggregateSession -> STM (Maybe (PublicKey,PublicKey))
+compatibleKeys c = do
+    isclosed <- isClosedTMChan (contactChannel c)
+    imap <- readTVar (contactSession c)
+    return $ case IntMap.elems imap of
+                    _ | isclosed -> Nothing -- none.
+                    con:_        -> Just ( toPublic (sOurKey $ singleSession con)
+                                         , sTheirUserKey (singleSession con))
+                    []           -> Nothing -- any.