Expire old tox key announcements after 5 minutes.

1 files changed, 30 insertions, 3 deletions

diff --git a/src/Network/Tox/Onion/Handlers.hs b/src/Network/Tox/Onion/Handlers.hs
index 99ef3c69..263d60bd 100644
--- a/src/Network/Tox/Onion/Handlers.hs
+++ b/src/Network/Tox/Onion/Handlers.hs
@@ -11,7 +11,7 @@ import Network.QueryResponse as QR hiding (Client)
 import qualified Network.QueryResponse as QR (Client)
 import Crypto.Tox
 import qualified Data.Wrapper.PSQ             as PSQ
-         ;import Data.Wrapper.PSQ             (PSQ)
+         ;import Data.Wrapper.PSQ             (PSQ,pattern (:->))
 #ifdef CRYPTONITE_BACKPORT
 import Crypto.Error.Types                     (CryptoFailable (..),
                                                throwCryptoError)
@@ -35,6 +35,12 @@ import Network.BitTorrent.DHT.Token           as Token
 
 import Control.Exception hiding (Handler)
 import Control.Monad
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent
+import GHC.Conc (labelThread)
+#endif
 import Control.Concurrent.STM
 import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
 import Network.Socket
@@ -159,7 +165,7 @@ toOnionDestination (AnnouncedRoute ni rpath) = OnionToOwner ni rpath
 -- structure.
 --
 data AnnouncedKeys = AnnouncedKeys
-    { keyByAge :: !(PSQ NodeId (Down POSIXTime{-Time at which they announced to you-})) -- TODO: timeout of 300 seconds
+    { keyByAge :: !(PSQ NodeId (POSIXTime{-Time at which they announced to you-}))
     , keyAssoc :: !(MinMaxPSQ' NodeId NodeDistance (Int{-count of route usage-},AnnouncedRoute))
         -- ^ PSQ using NodeId(user/public key) as Key
         --   and using 'NodeDistance' as priority.
@@ -172,12 +178,33 @@ data AnnouncedKeys = AnnouncedKeys
 
 insertKey :: POSIXTime -> NodeId -> AnnouncedRoute -> NodeDistance -> AnnouncedKeys -> AnnouncedKeys
 insertKey tm pub toxpath d keydb = AnnouncedKeys
-    { keyByAge = PSQ.insert pub (Down tm) (keyByAge keydb)
+    { keyByAge = PSQ.insert pub tm (keyByAge keydb)
     , keyAssoc = case MinMaxPSQ.lookup' pub (keyAssoc keydb) of
                     Just (_,(cnt,_)) -> MinMaxPSQ.insert' pub (cnt,toxpath) d (keyAssoc keydb)
                     Nothing          -> MinMaxPSQ.insert' pub (0  ,toxpath) d (keyAssoc keydb)
     }
 
+-- | Forks a thread to garbage-collect old key announcements.  Keys may be
+-- discarded after 5 minutes.
+forkAnnouncedKeysGC :: TVar AnnouncedKeys -> IO ThreadId
+forkAnnouncedKeysGC db = forkIO $ do
+    myThreadId >>= flip labelThread "gc:toxids"
+    fix $ \loop -> do
+    cutoff <- getPOSIXTime
+    threadDelay 300000000 -- 300 seconds
+    join $ atomically $ do
+        fix $ \gc -> do
+        keys <- readTVar db
+        case PSQ.minView (keyByAge keys) of
+            Nothing -> return loop
+            Just (pub :-> tm,kba')
+                | tm > cutoff -> return loop
+                | otherwise   -> do writeTVar db keys
+                                        { keyByAge = kba'
+                                        , keyAssoc = MinMaxPSQ.delete pub (keyAssoc keys)
+                                        }
+                                    gc
+
 areq :: Message -> Either String AnnounceRequest
 areq (OnionAnnounce asymm) = Right $ fst $ runIdentity $ asymmData asymm
 areq _                     = Left "Unexpected non-announce OnionMessage"