Avoid duplicate connection sessions in relay server.

1 files changed, 21 insertions, 12 deletions

diff --git a/dht/src/Network/Tox/Relay.hs b/dht/src/Network/Tox/Relay.hs
index 96838688..66ab4b71 100644
--- a/dht/src/Network/Tox/Relay.hs
+++ b/dht/src/Network/Tox/Relay.hs
@@ -24,7 +24,9 @@ import System.Timeout
 import Crypto.Tox
 import qualified Data.IntervalSet as IntSet
          ;import Data.IntervalSet (IntSet)
+import Data.Time.Clock.POSIX
 import Data.Tox.Relay
+import qualified Data.Wrapper.PSQInt as Int
 import Network.Address (getBindAddress)
 import Network.SocketLike
 import Network.StreamServer
@@ -47,16 +49,20 @@ hGetSized h = runGet get <$> B.hGet h len -- We treat parse-fail the same as EOF
 
 -- This type manages ConId assignments.
 data RelaySession = RelaySession
-    { indexPool  :: IntSet                                   -- ^ Ints that are either solicited or associated.
-    , solicited  :: Map PublicKey Int                        -- ^ Reserved ids, not yet in associated.
+    { indexPool  :: IntSet                                   -- ^ Ints are assigned.
+    , assigned   :: Map PublicKey Int                        -- ^ Assignments
     , associated :: IntMap ((ConId -> RelayPacket) -> IO ()) -- ^ Peers this session is connected to.
+    -- TODO: Timestamp PSQ for reclaiming indices.  The Bool will indicate
+    -- whether the index was ever associated.
+    , timestamps :: Int.PSQ (Bool,POSIXTime)
     }
 
 freshSession :: RelaySession
 freshSession = RelaySession
     { indexPool  = IntSet.empty
-    , solicited  = Map.empty
+    , assigned   = Map.empty
     , associated = IntMap.empty
+    , timestamps = Int.empty
     }
 
 disconnect :: TVar (Map PublicKey (RelayPacket -> IO (),TVar RelaySession))
@@ -174,31 +180,34 @@ handlePacket :: TVar (Map PublicKey (RelayPacket -> IO (), TVar RelaySession)) -
 handlePacket cons thistcp thisKey crypto sendOnion sendToClient session = \case
     RoutingRequest them -> join $ atomically $ do
         mySession <- readTVar session
-        mi <- case Map.lookup them (solicited mySession) of
+        -- TODO: Do we already have an association?
+        mi <- case Map.lookup them (assigned  mySession) of
             Nothing -> fmap join $ forM (IntSet.nearestOutsider 0 (indexPool mySession)) $ \i -> do
                 if -120 <= i && i <= 119
                     then do
                         writeTVar session mySession
                             { indexPool = IntSet.insert i (indexPool mySession)
-                            , solicited = Map.insert them i (solicited mySession)
+                            , assigned  = Map.insert them i (assigned  mySession)
                             }
                         return $ Just i
-                    else return Nothing -- No more slots available.
+                    else
+                        -- TODO: Attempt to reclaim an assigned but not yet associated connection.
+                        return Nothing -- No more slots available.
             Just i  -> return $ Just i
         notifyConnect <- fmap (join . join) $ forM mi $ \i -> do
             mp <- Map.lookup them <$> readTVar cons
             forM mp $ \(sendToThem,peer) -> do
                 theirSession <- readTVar peer
-                forM (Map.lookup thisKey $ solicited theirSession) $ \reserved_id -> do
+                forM (Map.lookup thisKey $ assigned  theirSession) $ \reserved_id -> do
                     let sendToThem'   f = sendToThem $ f $ key2c reserved_id
                         sendToClient' f = sendToClient $ f $ key2c i
                     writeTVar peer theirSession
-                        { solicited  = Map.delete thisKey (solicited theirSession)
-                        , associated = IntMap.insert reserved_id sendToClient' (associated theirSession)
+                        { -- assigned   = Map.insert thisKey reserved_id (assigned theirSession)
+                          associated = IntMap.insert reserved_id sendToClient' (associated theirSession)
                         }
                     writeTVar session mySession
-                        { solicited  = Map.delete them (solicited mySession)
-                        , associated = IntMap.insert i sendToThem' (associated mySession)
+                        { -- assigned   = Map.insert them i (assigned mySession)
+                          associated = IntMap.insert i sendToThem' (associated mySession)
                         }
                     return $ do
                         let showSession n k = "("++ show (key2c n) ++ ")" ++ showKey256 k
@@ -262,7 +271,7 @@ tcpRelay crypto udp_addr sendOnion = do
     b3389  <- getBindAddress "3389"  True
     b33445 <- getBindAddress "33445" True
     bany   <- getBindAddress ""      True
-    h <- streamServer ServerConfig
+    h <- forkStreamServer ServerConfig
                         { serverWarn    = dput XOnion
                         , serverSession = relaySession crypto clients cons sendOnion
                         }