1 files changed, 22 insertions, 69 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index f96ba707..d94f028f 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -35,9 +35,10 @@ module Network.BitTorrent.DHT.Session
        , routingInfo
        , routableAddress
        , getTimestamp
-       , SessionTokens
-       , sessionTokens
-       , contactInfo
+       -- , SessionTokens
+       -- , sessionTokens
+       -- , contactInfo
+       , dhtData
        , PeerStore
        , manager
 
@@ -55,8 +56,8 @@ module Network.BitTorrent.DHT.Session
        , runDHT
 
          -- ** Tokens
-       , grantToken
-       , checkToken
+       -- , grantToken
+       -- , checkToken
 
          -- ** Routing table
        , getTable
@@ -68,6 +69,7 @@ module Network.BitTorrent.DHT.Session
        , insertPeer
        , getPeerList
        , getPeerList1
+       , lookupPeers
        , insertTopic
        , deleteTopic
        , getSwarms
@@ -113,6 +115,7 @@ import Data.Time.Clock.POSIX
 import Data.Text as Text
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import Data.Serialize as S
+import Network.DHT.Types
 
 
 import Data.Torrent as Torrent
@@ -228,33 +231,6 @@ instance Default Options where
 
 seconds :: NominalDiffTime -> Int
 seconds dt = fromEnum (realToFrac dt :: Uni)
-
-{-----------------------------------------------------------------------
--- Tokens policy
------------------------------------------------------------------------}
-
-data SessionTokens = SessionTokens
-  { tokenMap    :: !TokenMap
-  , lastUpdate  :: !UTCTime
-  , maxInterval :: !NominalDiffTime
-  }
-
-nullSessionTokens :: IO SessionTokens
-nullSessionTokens = SessionTokens
-  <$> (tokens <$> liftIO randomIO)
-  <*> liftIO getCurrentTime
-  <*> pure defaultUpdateInterval
-
--- TODO invalidate *twice* if needed
-invalidateTokens :: UTCTime -> SessionTokens -> SessionTokens
-invalidateTokens curTime ts @ SessionTokens {..}
-  | curTime `diffUTCTime` lastUpdate > maxInterval = SessionTokens
-    { tokenMap    = update tokenMap
-    , lastUpdate  = curTime
-    , maxInterval = maxInterval
-    }
-  |                  otherwise                     = ts
-
 {-----------------------------------------------------------------------
 -- Session
 -----------------------------------------------------------------------}
@@ -277,9 +253,8 @@ data Node raw dht u ip = Node
   , resources     :: !InternalState
   , manager       :: !(Manager raw dht) -- ^ RPC manager;
   , routingInfo   :: !(TVar    (Maybe (R.Info dht ip u))) -- ^ search table;
-  , contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
   , announceInfo  :: !(TVar     AnnounceSet       ) -- ^ to publish by this node;
-  , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
+  , dhtData       :: DHTData dht ip
   , loggerFun     :: !LogFun
   }
 
@@ -371,6 +346,7 @@ locFromCS cs = case getCallStack cs of
 newNode :: ( Address ip
            , FiniteBits (NodeId dht)
            , Serialize (NodeId dht)
+           , Kademlia dht
            )
           => -- [NodeHandler] -- ^ handlers to run on accepted queries;
              Options       -- ^ various dht options;
@@ -389,12 +365,12 @@ newNode opts naddr logger mbid = do
       s      <- getInternalState
       (_, m) <- allocate (newManager rpcOpts (logt logger) nodeAddr []) closeManager
       liftIO $ do
+        dta <- initializeDHTData
         myId   <- maybe genNodeId return mbid
         node   <- Node opts myId s m
                 <$> atomically (newTVar Nothing)
-                <*> newTVarIO def
                 <*> newTVarIO S.empty
-                <*> (newTVarIO =<< nullSessionTokens)
+                <*> pure dta
                 <*> pure logger
         return node
 
@@ -415,29 +391,6 @@ runDHT node action = runReaderT (unDHT action) node
 -- /pick a random ID/ in the range of the bucket and perform a
 -- find_nodes search on it.
 
-{-----------------------------------------------------------------------
--- Tokens
------------------------------------------------------------------------}
-
-tryUpdateSecret :: TVar SessionTokens -> IO ()
-tryUpdateSecret toks = do
-  curTime <- liftIO getCurrentTime
-  liftIO $ atomically $ modifyTVar' toks (invalidateTokens curTime)
-
-grantToken :: Hashable a => TVar SessionTokens -> NodeAddr a -> IO Token
-grantToken sessionTokens addr = do
-  tryUpdateSecret sessionTokens
-  toks <- readTVarIO sessionTokens
-  return $ T.lookup addr $ tokenMap toks
-
--- | Throws 'HandlerError' if the token is invalid or already
--- expired. See 'TokenMap' for details.
-checkToken :: Hashable a => TVar SessionTokens -> NodeAddr a -> Token -> IO Bool
-checkToken sessionTokens addr questionableToken = do
-  tryUpdateSecret sessionTokens
-  toks <- readTVarIO sessionTokens
-  return $ T.member addr questionableToken (tokenMap toks)
-
 
 {-----------------------------------------------------------------------
 -- Routing table
@@ -475,28 +428,28 @@ getTable = do
   let nil = nullTable myId (optBucketCount opts)
   liftIO (maybe nil R.myBuckets <$> atomically (readTVar var))
 
-getSwarms :: Ord ip => DHT raw dht u ip [ (InfoHash, Int, Maybe ByteString) ]
+getSwarms :: Ord ip => DHT raw KMessageOf u ip [ (InfoHash, Int, Maybe ByteString) ]
 getSwarms = do
-    store <- asks contactInfo >>= liftIO . atomically . readTVar
+    store <- asks (contactInfo . dhtData) >>= liftIO . atomically . readTVar
     return $ P.knownSwarms store
 
-savePeerStore :: (Ord ip, Address ip) => DHT raw dht u ip ByteString
+savePeerStore :: (Ord ip, Address ip) => DHT raw KMessageOf u ip ByteString
 savePeerStore = do
-    var <- asks contactInfo
+    var <- asks (contactInfo . dhtData)
     peers <- liftIO $ atomically $ readTVar var
     return $ S.encode peers
 
-mergeSavedPeers :: (Ord ip, Address ip) => ByteString -> DHT raw dht u ip ()
+mergeSavedPeers :: (Ord ip, Address ip) => ByteString -> DHT raw KMessageOf u ip ()
 mergeSavedPeers bs = do
-    var <- asks contactInfo
+    var <- asks (contactInfo . dhtData)
     case S.decode bs of
         Right newbies -> liftIO $ atomically $ modifyTVar' var (<> newbies)
         Left _        -> return ()
 
 
-allPeers :: Ord ip => InfoHash -> DHT raw dht u ip [ PeerAddr ip ]
+allPeers :: Ord ip => InfoHash -> DHT raw KMessageOf u ip [ PeerAddr ip ]
 allPeers ih = do
-    store <- asks contactInfo >>= liftIO . atomically . readTVar
+    store <- asks (contactInfo . dhtData) >>= liftIO . atomically . readTVar
     return $ P.lookup ih store
 
 -- | Find a set of closest nodes from routing table of this node. (in
@@ -566,7 +519,7 @@ getTimestamp = do
 --
 getPeerList :: Ord ip => InfoHash -> DHT raw KMessageOf () ip (PeerList ip)
 getPeerList ih = do
-  var <- asks contactInfo
+  var <- asks (contactInfo . dhtData)
   ps <- liftIO $ lookupPeers var ih
   if L.null ps
     then Left <$> getClosest ih
@@ -574,7 +527,7 @@ getPeerList ih = do
 
 getPeerList1 :: Ord ip => DHT raw KMessageOf () ip (InfoHash -> IO (PeerList ip))
 getPeerList1 = do
-  var <- asks contactInfo
+  var <- asks (contactInfo . dhtData)
   getclosest <- getClosest1
   return $ \ih -> do
       ps <- lookupPeers var ih