2 files changed, 100 insertions, 0 deletions

diff --git a/dht/dht-client.cabal b/dht/dht-client.cabal
index dc9b7ca2..19f4f2b3 100644
--- a/dht/dht-client.cabal
+++ b/dht/dht-client.cabal
@@ -99,6 +99,7 @@ library
                     Network.Tox.Handshake
                     Network.Tox.NodeId
                     Network.Tox.Avahi
+                    Network.Tox.RelayPinger
                     Network.UPNP
                     Network.QueryResponse.TCP
                     Network.Tox.Relay
diff --git a/dht/src/Network/Tox/RelayPinger.hs b/dht/src/Network/Tox/RelayPinger.hs
new file mode 100644
index 00000000..00c6f65a
--- /dev/null
+++ b/dht/src/Network/Tox/RelayPinger.hs
@@ -0,0 +1,99 @@
+{-# LANGUAGE NamedFieldPuns  #-}
+{-# LANGUAGE PatternSynonyms #-}
+module Network.Tox.RelayPinger where
+
+import Announcer
+import Control.Arrow
+import Control.Concurrent.STM
+import Control.Concurrent.ThreadUtil
+import Control.Monad
+import Data.Ord
+import qualified Data.Serialize      as S
+import Data.Time.Clock               (NominalDiffTime)
+import Data.Time.Clock.POSIX
+import qualified Data.Wrapper.PSQ    as PSQ
+         ;import Data.Wrapper.PSQ    (pattern (:->), PSQ)
+import Network.Kademlia.Routing
+import Network.QueryResponse
+import Network.Tox.TCP               as TCP
+
+data RelayPriority = RelayPriority
+    { tcpPortRank      :: Int
+    , xorDeltaDistance :: NodeId
+    }
+ deriving (Eq,Ord)
+
+data RelayPinger = RelayPinger
+    { relayAnnouncer     :: Announcer
+    , relayKademliaSpace :: KademliaSpace NodeId NodeInfo
+    , relayClient        :: RelayClient
+    , relayQueue         :: TVar (Int, PSQ NodeInfo (Down RelayPriority))
+    , relaySelf          :: NodeId
+    }
+
+forkRelayPinger :: KademliaSpace NodeId NodeInfo -> RelayClient -> IO RelayPinger
+forkRelayPinger kad cli = do
+    a <- forkAnnouncer
+    q <- newTVarIO (0,PSQ.empty)
+    self <- kademliaLocation kad <$> clientAddress cli Nothing
+    return RelayPinger
+        { relayAnnouncer     = a
+        , relayKademliaSpace = kad
+        , relayClient        = cli
+        , relayQueue         = q
+        , relaySelf          = self
+        }
+
+pinginterval :: NominalDiffTime
+pinginterval = 120.0 -- 2 minutes
+
+
+rankTCPPort :: NodeInfo -> Int
+rankTCPPort NodeInfo{tcpPort = 443}  = 1
+rankTCPPort NodeInfo{tcpPort = 80}   = 1
+rankTCPPort NodeInfo{tcpPort = 3389} = 2
+rankTCPPort _                        = 3
+
+addRelay :: RelayPinger -> TCP.NodeInfo -> STM ()
+addRelay (RelayPinger a kad cli que self) ni = do
+    let nid = kademliaLocation kad ni
+        nidk = decodeAnnounceKey a $ S.encode nid
+        pingit = ScheduledItem $ \a nidk tm -> do
+            scheduleRel a nidk pingit pinginterval -- Reschedule the next ping.
+            return $ do
+                _ <- tcpPing cli ni
+                -- TODO: Remove after ping timeout?
+                return ()
+    modifyTVar' que $ \(n,q) ->
+                    ( (,) $! (if PSQ.member ni q then n else n+1) )
+                    $ PSQ.insert
+                        ni
+                        (Down $ RelayPriority (rankTCPPort ni) (kademliaXor kad self nid))
+                        q
+    scheduleRel a nidk pingit pinginterval
+
+delRelay :: RelayPinger -> TCP.NodeInfo -> STM ()
+delRelay (RelayPinger a kad _ que _) ni = do
+    let nid = kademliaLocation kad ni
+        nidk = decodeAnnounceKey a $ S.encode nid
+    (n,q) <- readTVar que
+    writeTVar que $ ( (,) $! (if PSQ.member ni q then n - 1 else n) )
+                  $ PSQ.delete ni q
+    unschedule a nidk
+
+bumpRelay :: RelayPinger -> TCP.NodeInfo -> STM ()
+bumpRelay rp@(RelayPinger a kad cli que self) ni = do
+    let nid = kademliaLocation kad ni
+    (cnt,q) <- readTVar que
+    if cnt < 4
+        then addRelay rp ni
+        else case PSQ.minView q of
+            Nothing            -> addRelay rp ni
+            Just (r :-> Down p, q')
+                | let pnew = RelayPriority (rankTCPPort ni) (kademliaXor kad self nid)
+                , pnew < p
+                -> delRelay rp r >> addRelay rp ni
+            _ -> return ()
+
+currentRelays :: RelayPinger -> STM (Int,[TCP.NodeInfo])
+currentRelays rp = second PSQ.keys <$> readTVar (relayQueue rp)