1 files changed, 31 insertions, 12 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 5a8d64ef..44a5f0e9 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -29,8 +29,10 @@ module Network.BitTorrent.DHT.Session
          -- * Session
        , Node
        , options
-       , thisNodeId
-       , routingTable
+       , tentativeNodeId
+       , myNodeIdAccordingTo
+       , routingInfo
+       , routableAddress
 
          -- ** Initialization
        , LogFun
@@ -239,14 +241,14 @@ data Node ip = Node
 
     -- | Pseudo-unique self-assigned session identifier. This value is
     -- constant during DHT session and (optionally) between sessions.
-  , thisNodeId    :: !NodeId
+  , tentativeNodeId    :: !NodeId
 
   , resources     :: !InternalState
-  , manager       :: !(Manager (DHT       ip)) -- ^ RPC manager;
-  , routingTable  :: !(TVar    (Table     ip)) -- ^ search table;
-  , contactInfo   :: !(TVar    (PeerStore ip)) -- ^ published by other nodes;
-  , announceInfo  :: !(TVar     AnnounceSet  ) -- ^ to publish by this node;
-  , sessionTokens :: !(TVar     SessionTokens) -- ^ query session IDs.
+  , manager       :: !(Manager (DHT       ip     )) -- ^ RPC manager;
+  , routingInfo   :: !(TVar    (Maybe (R.Info ip))) -- ^ search table;
+  , contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
+  , announceInfo  :: !(TVar     AnnounceSet       ) -- ^ to publish by this node;
+  , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
   , loggerFun     :: !LogFun
   }
 
@@ -323,7 +325,7 @@ newNode hs opts naddr logger mbid = do
       liftIO $ do
         myId   <- maybe genNodeId return mbid
         node   <- Node opts myId s m
-                <$> atomically (newTVar (nullTable myId (optBucketCount opts)))
+                <$> atomically (newTVar Nothing)
                 <*> newTVarIO def
                 <*> newTVarIO S.empty
                 <*> (newTVarIO =<< nullSessionTokens)
@@ -372,16 +374,33 @@ checkToken addr questionableToken = do
   toks <- asks sessionTokens >>= liftIO . readTVarIO
   return $ T.member addr questionableToken (tokenMap toks)
 
+
 {-----------------------------------------------------------------------
 -- Routing table
 -----------------------------------------------------------------------}
 
+-- | This nodes externally routable address reported by remote peers.
+routableAddress :: DHT ip (Maybe ip)
+routableAddress = do
+    info <- asks routingInfo >>= liftIO . atomically . readTVar
+    return $ myAddress <$> info
+
+-- | The current NodeId that the given remote node should know us by.
+myNodeIdAccordingTo :: NodeAddr ip -> DHT ip NodeId
+myNodeIdAccordingTo _ = do
+  info <- asks routingInfo >>= liftIO . atomically . readTVar
+  fallback <- asks tentativeNodeId
+  return $ maybe fallback 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.
-getTable :: DHT ip (Table ip)
+getTable :: Eq ip => DHT ip (Table ip)
 getTable = do
-  var <- asks routingTable
-  liftIO (atomically $ readTVar var)
+  Node { tentativeNodeId = myId
+       , routingInfo = var
+       , options = opts } <- ask
+  let nil = nullTable myId (optBucketCount opts)
+  liftIO (maybe nil R.myBuckets <$> atomically (readTVar var))
 
 -- | Find a set of closest nodes from routing table of this node. (in
 -- no particular order)