3 files changed, 142 insertions, 82 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 8bc423a3..6d31eab2 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -76,8 +76,13 @@ import Data.Typeable
 import Data.Monoid
 import Network.DatagramServer.Mainline (KMessageOf)
 import qualified Network.DatagramServer as KRPC (listen, Protocol(..))
-
-
+import Network.DatagramServer.Types
+import Network.DHT.Types
+import Data.Bits
+import Data.Default
+import Text.PrettyPrint.HughesPJClass hiding (($$), (<>))
+import Network.KRPC.Method
+import Network.BitTorrent.DHT.Query (DataHandlers)
 
 {-----------------------------------------------------------------------
 --  DHT types
@@ -96,7 +101,31 @@ fullLogging :: LogSource -> LogLevel -> Bool
 fullLogging _ _ = True
 
 -- | Run DHT on specified port. <add note about resources>
-dht :: (Ord ip, Address ip)
+dht ::
+    ( Ord ip
+    , Address ip
+    , Functor dht
+    , Ord (NodeId dht)
+    , FiniteBits (NodeId dht)
+    , Serialize (NodeId dht)
+    , Show (NodeId dht)
+    , SerializableTo raw (Response dht (Ping dht))
+    , SerializableTo raw (Query dht (Ping dht))
+    , SerializableTo raw (Response dht (NodeFound dht ip))
+    , SerializableTo raw (Query dht (FindNode dht ip))
+    , Ord (TransactionID dht)
+    , Serialize (TransactionID dht)
+    , Eq (QueryMethod dht)
+    , Show (QueryMethod dht)
+    , Pretty (NodeInfo dht ip u)
+    , Kademlia dht
+    , KRPC (Query dht (FindNode dht ip)) (Response dht (NodeFound dht ip))
+    , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+    , DataHandlers raw dht
+    , WireFormat raw dht
+    , Show u
+    , Default u
+    )
     => Options     -- ^ normally you need to use 'Data.Default.def';
     -> NodeAddr ip -- ^ address to bind this node;
     -> (LogSource -> LogLevel -> Bool) -- ^ use 'fullLogging' as a noisy default
@@ -179,7 +208,33 @@ resolveHostName NodeAddr {..} = do
 --
 --   This operation do block, use
 --   'Control.Concurrent.Async.Lifted.async' if needed.
-bootstrap :: forall raw dht u ip. Address ip => Maybe BS.ByteString -> [NodeAddr ip] -> DHT raw dht u ip ()
+bootstrap :: forall raw dht u ip.
+    ( Ord ip
+    , Address ip
+    , Functor dht
+    , Ord (NodeId dht)
+    , FiniteBits (NodeId dht)
+    , Serialize (NodeId dht)
+    , Show (NodeId dht)
+    , Pretty (NodeId dht)
+    , SerializableTo raw (Response dht (Ping dht))
+    , SerializableTo raw (Query dht (Ping dht))
+    , SerializableTo raw (Response dht (NodeFound dht ip))
+    , SerializableTo raw (Query dht (FindNode dht ip))
+    , Ord (TransactionID dht)
+    , Serialize (TransactionID dht)
+    , Eq (QueryMethod dht)
+    , Show (QueryMethod dht)
+    , Pretty (NodeInfo dht ip u)
+    , Kademlia dht
+    , KRPC (Query dht (FindNode dht ip)) (Response dht (NodeFound dht ip))
+    , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+    , DataHandlers raw dht
+    , WireFormat raw dht
+    , Show u
+    , Default u
+    , Serialize u
+    ) => Maybe BS.ByteString -> [NodeAddr ip] -> DHT raw dht u ip ()
 bootstrap mbs startNodes = do
   restored <-
     case decode <$> mbs of
@@ -192,7 +247,7 @@ bootstrap mbs startNodes = do
   let searchAll aliveNodes = do
         nid <- myNodeIdAccordingTo (error "FIXME")
         ns <- bgsearch ioFindNodes nid
-        return ( ns :: [NodeInfo KMessageOf ip ()] )
+        return ( ns :: [NodeInfo dht ip u] )
   input_nodes <- (restored ++) . T.toList <$> getTable
   -- Step 1: Use iterative searches to flesh out the table..
   do let knowns = map (map $ nodeAddr . fst) input_nodes
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index e5d9bd5f..67dc4541 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -16,6 +16,8 @@
 {-# LANGUAGE TupleSections       #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 module Network.BitTorrent.DHT.Query
        ( -- * Handler
          -- | To bind specific set of handlers you need to pass
@@ -25,6 +27,7 @@ module Network.BitTorrent.DHT.Query
        , getPeersH
        , announceH
        , defaultHandlers
+       , DataHandlers
 
          -- * Query
          -- ** Basic
@@ -113,6 +116,7 @@ import Data.Serialize
 import System.IO.Unsafe (unsafeInterleaveIO)
 import Data.String
 
+
 {-----------------------------------------------------------------------
 --  Handlers
 -----------------------------------------------------------------------}
@@ -215,20 +219,68 @@ kademliaHandlers logger = do
            , handler (nameFindNodes dht) $ findNodeH getclosest
            ]
 
+class DataHandlers raw dht where
+    dataHandlers ::
+        ( Ord ip , Hashable ip, Typeable ip, Serialize ip) =>
+         (NodeId dht  -> IO [NodeInfo dht ip ()])
+        -> DHTData dht ip
+        -> [MethodHandler raw dht ip]
+
+instance DataHandlers BValue KMessageOf where
+    dataHandlers = bthandlers
+
+bthandlers ::
+    ( Ord ip , Hashable ip, Typeable ip, Serialize ip) =>
+     (NodeId KMessageOf  -> IO [NodeInfo KMessageOf ip ()])
+    -> DHTData KMessageOf ip
+    -> [MethodHandler BValue KMessageOf ip]
+bthandlers getclosest dta =
+    [ MethodHandler "get_peers" $ getPeersH (getpeers dta) (sessionTokens dta)
+    , MethodHandler "announce_peer" $ announceH (contactInfo dta) (sessionTokens dta)
+    ]
+ where
+    getpeers dta ih = do
+      ps <- lookupPeers (contactInfo dta) ih
+      if L.null ps
+        then Left <$> getclosest (toNodeId ih)
+        else return (Right ps)
+
+data MethodHandler raw dht ip =
+        forall a b. ( SerializableTo raw (Response dht b)
+                    , SerializableTo raw (Query dht a)
+                    ) => MethodHandler (QueryMethod dht) (NodeAddr ip -> a -> IO b)
 
 -- | Includes all default query handlers.
-defaultHandlers :: forall ip. (Eq ip, Ord ip, Address ip) => LogFun -> DHT BValue KMessageOf () ip [NodeHandler]
+defaultHandlers :: forall raw dht u ip.
+                    ( Ord (TransactionID dht)
+                    , Ord (NodeId dht)
+                    , Show u
+                    , SerializableTo raw (Response dht (Ping dht))
+                    , SerializableTo raw (Query dht (Ping dht))
+                    , Show (QueryMethod dht)
+                    , Show (NodeId dht)
+                    , FiniteBits (NodeId dht)
+                    , Default u
+                    , Serialize (TransactionID dht)
+                    , WireFormat raw dht
+                    , Kademlia dht
+                    , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+                    , Functor dht
+                    , Pretty (NodeInfo dht ip u)
+                    , KRPC (Query dht (FindNode dht ip)) (Response dht (NodeFound dht ip))
+                    , SerializableTo raw (Response dht (NodeFound dht ip))
+                    , SerializableTo raw (Query dht (FindNode dht ip))
+                    , Eq ip, Ord ip, Address ip, DataHandlers raw dht
+                    ) => LogFun -> DHT raw dht u ip [Handler IO dht raw]
 defaultHandlers logger = do
     groknode <- insertNode1
     mynid <- myNodeIdAccordingTo1
-    let handler :: KRPC (Query KMessageOf a) (Response KMessageOf b) => QueryMethod KMessageOf -> (NodeAddr ip -> a -> IO b) -> NodeHandler
-        handler = nodeHandler groknode mynid (logt logger)
-    toks <- asks sessionTokens
-    peers <- asks contactInfo
-    getpeers <- getPeerList1
+    let handler :: MethodHandler raw dht ip -> Handler IO dht raw
+        handler (MethodHandler name action) = nodeHandler groknode mynid (logt logger) name action
+    dta <- asks dhtData
+    getclosest <- getClosest1
     hs <- kademliaHandlers logger
-    return $ hs ++ [ handler "get_peers"     $ getPeersH getpeers toks
-                   , handler "announce_peer" $ announceH peers toks ]
+    return $ hs ++ L.map handler (dataHandlers (fmap (fmap (fmap (const ()))) . getclosest) dta)
 
 {-----------------------------------------------------------------------
 --  Basic queries
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