1 files changed, 70 insertions, 32 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 7e87df6c..d8665773 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -31,12 +31,19 @@ module Network.BitTorrent.DHT.Session
        , options
        , tentativeNodeId
        , myNodeIdAccordingTo
+       , myNodeIdAccordingTo1
        , routingInfo
        , routableAddress
        , getTimestamp
+       , SessionTokens
+       , sessionTokens
+       , contactInfo
+       , PeerStore
+       , manager
 
          -- ** Initialization
        , LogFun
+       , logt
        , NodeHandler
        , newNode
        , closeNode
@@ -54,11 +61,13 @@ module Network.BitTorrent.DHT.Session
          -- ** Routing table
        , getTable
        , getClosest
+       , getClosest1
 
 #ifdef VERSION_bencoding
          -- ** Peer storage
        , insertPeer
        , getPeerList
+       , getPeerList1
        , insertTopic
        , deleteTopic
        , getSwarms
@@ -333,7 +342,7 @@ instance MonadLogger (DHT ip) where
     liftIO $ logger loc src lvl (toLogStr msg)
 
 #ifdef VERSION_bencoding
-type NodeHandler ip = Handler IO KMessageOf BValue
+type NodeHandler = Handler IO KMessageOf BValue
 #else
 type NodeHandler ip = Handler (DHT ip) Tox.Message ByteString
 #endif
@@ -368,10 +377,10 @@ locFromCS cs = case getCallStack cs of
 -- 'closeNode' function, otherwise socket or other scarce resources may
 -- leak.
 newNode :: Address ip
-          => [NodeHandler ip] -- ^ handlers to run on accepted queries;
-          -> Options          -- ^ various dht options;
-          -> NodeAddr ip      -- ^ node address to bind;
-          -> LogFun           -- ^
+          => [NodeHandler] -- ^ handlers to run on accepted queries;
+          -> Options       -- ^ various dht options;
+          -> NodeAddr ip   -- ^ node address to bind;
+          -> LogFun        -- ^ invoked on log messages;
 #ifdef VERSION_bencoding
           -> Maybe (NodeId KMessageOf)     -- ^ use this NodeId, if not given a new one is generated.
 #else
@@ -420,24 +429,23 @@ runDHT node action = runReaderT (unDHT action) node
 -- Tokens
 -----------------------------------------------------------------------}
 
-tryUpdateSecret :: DHT ip ()
-tryUpdateSecret = do
+tryUpdateSecret :: TVar SessionTokens -> IO ()
+tryUpdateSecret toks = do
   curTime <- liftIO getCurrentTime
-  toks    <- asks sessionTokens
   liftIO $ atomically $ modifyTVar' toks (invalidateTokens curTime)
 
-grantToken :: Hashable a => NodeAddr a -> DHT ip Token
-grantToken addr = do
-  tryUpdateSecret
-  toks <- asks sessionTokens >>= liftIO . readTVarIO
+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 => NodeAddr a -> Token -> DHT ip Bool
-checkToken addr questionableToken = do
-  tryUpdateSecret
-  toks <- asks sessionTokens >>= liftIO . readTVarIO
+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)
 
 
@@ -463,6 +471,14 @@ myNodeIdAccordingTo _ = do
         (return . myNodeId)
         info
 
+myNodeIdAccordingTo1 :: DHT ip ( NodeAddr ip -> IO (NodeId KMessageOf) )
+myNodeIdAccordingTo1 = do
+  var <- asks routingInfo
+  tid <- asks tentativeNodeId
+  return $ \ _ -> do
+      info <- atomically $ readTVar var
+      return $ maybe tid myNodeId info
+
 -- | Get current routing table. Normally you don't need to use this
 -- function, but it can be usefull for debugging and profiling purposes.
 #ifdef VERSION_bencoding
@@ -515,38 +531,48 @@ getClosest node = do
   k <- asks (optK . options)
   kclosest k node <$> getTable
 
+getClosest1 :: ( Eq ip
+               , TableKey KMessageOf k
+               ) => DHT ip (k -> IO [NodeInfo KMessageOf ip ()])
+getClosest1 = do
+    k <- asks (optK . options)
+    nobkts <- asks (optBucketCount . options)
+    myid <- asks tentativeNodeId
+    var <- asks routingInfo
+    return $ \node -> do nfo <- atomically $ readTVar var
+                         let tbl = maybe (nullTable myid nobkts) R.myBuckets nfo
+                         return $ kclosest k node tbl
+
 {-----------------------------------------------------------------------
 -- Peer storage
 -----------------------------------------------------------------------}
 
-refreshContacts :: DHT ip ()
-refreshContacts =
+refreshContacts :: Ord ip => TVar (PeerStore ip) -> IO ()
+refreshContacts var =
     -- TODO limit dht peer store in size (probably by removing oldest peers)
     return ()
 
 
 -- | Insert peer to peer store. Used to handle announce requests.
-insertPeer :: Ord ip => InfoHash -> Maybe ByteString -> PeerAddr ip -> DHT ip ()
-insertPeer ih name addr = do
-  refreshContacts
-  var <- asks contactInfo
-  liftIO $ atomically $ modifyTVar' var (P.insertPeer ih name addr)
+insertPeer :: Ord ip => TVar (PeerStore ip) -> InfoHash -> Maybe ByteString -> PeerAddr ip -> IO ()
+insertPeer var ih name addr = do
+  refreshContacts var
+  atomically $ modifyTVar' var (P.insertPeer ih name addr)
 
 -- | Get peer set for specific swarm.
-lookupPeers :: Ord ip => InfoHash -> DHT ip [PeerAddr ip]
-lookupPeers ih = do
-  refreshContacts
-  var <- asks contactInfo
+lookupPeers :: Ord ip => TVar (PeerStore ip) -> InfoHash -> IO [PeerAddr ip]
+lookupPeers var ih = do
+  refreshContacts var
   tm <- getTimestamp
-  liftIO $ atomically $ do
+  atomically $ do
     (ps,store') <- P.freshPeers ih tm <$> readTVar var
     writeTVar var store'
     return ps
 
-getTimestamp :: DHT ip Timestamp
+getTimestamp :: IO Timestamp
 getTimestamp = do
-  utcTime <- liftIO $ getCurrentTime
-  $(logDebugS) "routing.make_timestamp" (Text.pack (render (pPrint utcTime)))
+  utcTime <- getCurrentTime
+  -- $(logDebugS) "routing.make_timestamp" (Text.pack (render (pPrint utcTime)))
   return $ utcTimeToPOSIXSeconds utcTime
 
 
@@ -557,11 +583,23 @@ getTimestamp = do
 --
 getPeerList :: Ord ip => InfoHash -> DHT ip (PeerList ip)
 getPeerList ih = do
-  ps <- lookupPeers ih
+  var <- asks contactInfo
+  ps <- liftIO $ lookupPeers var ih
   if L.null ps
     then Left <$> getClosest ih
     else return (Right ps)
 
+getPeerList1 :: Ord ip => DHT ip (InfoHash -> IO (PeerList ip))
+getPeerList1 = do
+  var <- asks contactInfo
+  getclosest <- getClosest1
+  return $ \ih -> do
+      ps <- lookupPeers var ih
+      if L.null ps
+        then Left <$> getclosest ih
+        else return (Right ps)
+
+
 insertTopic :: InfoHash -> PortNumber -> DHT ip ()
 insertTopic ih p = do
   var <- asks announceInfo