8 files changed, 418 insertions, 123 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index a358a2b8..c46ed17f 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -78,6 +78,7 @@ library
   hs-source-dirs:      src
   exposed-modules:     Network.DatagramServer
                        Network.DatagramServer.Mainline
+                       Network.DatagramServer.Tox
                        Network.DatagramServer.Types
                        Network.DatagramServer.Error
                        Network.DHT
@@ -392,6 +393,8 @@ executable dhtd
                      , unix
                      , containers
                      , stm
+                     , cereal
+                     , bencoding
   if flag(thread-debug)
     build-depends: time
     cpp-options:   -DTHREAD_DEBUG
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 8496dd5a..99ff7218 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE NondecreasingIndentation #-}
 {-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE FlexibleInstances     #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE LambdaCase            #-}
 {-# LANGUAGE OverloadedStrings     #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
@@ -7,6 +9,7 @@
 {-# LANGUAGE RecordWildCards       #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE RankNTypes #-}
 
 import Control.Arrow
 import Control.Monad
@@ -39,6 +42,7 @@ import Network.BitTorrent.DHT.Query
 import Network.DHT.Mainline (FindNode(..),NodeFound(..),GetPeers(..),GotPeers(..))
 import Network.DatagramServer (QueryFailure(..))
 import Network.DatagramServer.Mainline (ReflectedIP(..),KMessageOf)
+import qualified Network.DatagramServer.Tox as Tox
 import qualified Network.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 import Network.SocketLike
@@ -53,6 +57,14 @@ import Control.Concurrent
 #endif
 import Control.Concurrent.STM
 import System.Environment
+import Data.BEncode (BValue)
+import Network.DHT.Types
+import Network.DatagramServer.Types
+import Data.Bits
+import Data.Serialize
+import Network.KRPC.Method
+import Data.Typeable
+import GHC.Generics
 
 mkNodeAddr :: SockAddr -> NodeAddr IPv4
 mkNodeAddr addr = NodeAddr (fromJust $ fromSockAddr addr)
@@ -72,20 +84,27 @@ showReport kvs = do
     (k,v) <- kvs
     concat [ printf "  %-*s" (colwidth+1) k, v, "\n" ]
 
-showEnry :: Show a => (NodeInfo KMessageOf a (), t) -> [Char]
+showEnry ::
+    ( Show a
+    , Pretty (NodeId dht)
+    ) => (NodeInfo dht a u, t) -> [Char]
 showEnry (n,_) = intercalate " "
     [ show $ pPrint (nodeId n)
     , show $ nodeAddr n
     ]
 
-printTable :: DHT IPv4 ()
+printTable ::
+    ( Pretty (NodeId dht)
+    ) => DHT raw dht u IPv4 ()
 printTable = do
     t <- showTable
     liftIO $ do
         putStrLn t
         hFlush stdout
 
-showTable :: DHT IPv4 String
+showTable ::
+    ( Pretty (NodeId dht)
+    ) => DHT raw dht u IPv4 String
 showTable = do
     nodes <- R.toList <$> getTable
     return $ showReport
@@ -109,7 +128,7 @@ noLogging _ _ = False
 allNoise :: LogSource -> LogLevel ->  Bool
 allNoise _ _ = True
 
-resume :: DHT IPv4 (Maybe B.ByteString)
+resume :: DHT raw dht u IPv4 (Maybe B.ByteString)
 resume = do
     restore_attempt <- liftIO $ tryIOError $ B.readFile "dht-nodes.dat"
     saved_nodes <-
@@ -119,7 +138,30 @@ resume = do
             restore_attempt
     return saved_nodes
 
-godht :: String -> (NodeAddr IPv4 -> NodeId _ -> DHT IPv4 b) -> IO b
+godht ::
+    ( Eq (QueryMethod dht)
+    , Show (QueryMethod dht)
+    , Functor dht
+    , Ord (TransactionID dht)
+    , Serialize (TransactionID dht)
+    , Kademlia dht
+    , WireFormat raw dht
+    , DataHandlers raw dht
+    , SerializableTo raw (Query dht (FindNode dht IPv4))
+    , SerializableTo raw (Response dht (NodeFound dht IPv4))
+    , SerializableTo raw (Query dht (Ping dht))
+    , SerializableTo raw (Response dht (Ping dht))
+    , KRPC (Query dht (FindNode dht IPv4)) (Response dht (NodeFound dht IPv4))
+    , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+    , Ord (NodeId dht)
+    , FiniteBits (NodeId dht)
+    , Serialize (NodeId dht)
+    , Show (NodeId dht)
+    , Pretty (NodeId dht)
+    , Pretty (NodeInfo dht IPv4 u)
+    , Default u
+    , Show u
+    ) => String -> (NodeAddr IPv4 -> NodeId dht -> DHT raw dht u IPv4 b) -> IO b
 godht p f = do
     a <- btBindAddr p False
     dht def { optTimeout = 5 } a noDebugPrints $ do
@@ -138,21 +180,118 @@ hPutClient h s = hPutStr h ('.' : marshalForClient s)
 hPutClientChunk :: Handle -> String -> IO ()
 hPutClientChunk h s = hPutStr h (' ' : marshalForClient s)
 
-clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO ()
-clientSession st signalQuit sock n h = do
+data GenericDHT ip a
+    = GenericDHT
+        (forall raw dht u.
+            ( Eq (QueryMethod dht)
+            , Show (QueryMethod dht)
+            , Functor dht
+            , Ord (TransactionID dht)
+            , Serialize (TransactionID dht)
+            , Kademlia dht
+            , WireFormat raw dht
+            , DataHandlers raw dht
+            , SerializableTo raw (Query dht (FindNode dht ip))
+            , SerializableTo raw (Response dht (NodeFound dht ip))
+            , SerializableTo raw (Query dht (Ping dht))
+            , SerializableTo raw (Response dht (Ping dht))
+            , KRPC (Query dht (FindNode dht ip)) (Response dht (NodeFound dht ip))
+            , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+            , Ord (NodeId dht)
+            , FiniteBits (NodeId dht)
+            , Serialize (NodeId dht)
+            , Show (NodeId dht)
+            , Pretty (NodeId dht)
+            , Pretty (NodeInfo dht ip ())
+            , Pretty (NodeInfo dht ip u)
+            , Default u
+            , Show u
+            , Read (NodeId dht)
+            ) => DHT raw dht u ip a)
+    | BtDHT (DHT BValue KMessageOf () ip a)
+
+dhtType :: DHT raw dht u ip (Proxy dht)
+dhtType = return Proxy
+
+nodeIdType :: NodeId dht -> DHT raw dht u ip ()
+nodeIdType _ = return ()
+
+nodeAddrType :: NodeAddr ip -> DHT raw dht u ip ()
+nodeAddrType _ = return ()
+
+ipType :: f dht ip -> DHT raw dht u ip ()
+ipType _ = return ()
+
+instance Kademlia Tox.Message where
+instance Pretty (NodeId Tox.Message) where
+instance Pretty (NodeInfo Tox.Message IPv4 ()) where
+instance Pretty (NodeInfo Tox.Message IPv4 Bool) where -- TODO
+instance Read (NodeId KMessageOf) where -- TODO
+instance Read (NodeId Tox.Message) where -- TODO
+instance Serialize (FindNode Tox.Message IPv4) where
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (NodeFound Tox.Message IPv4) where
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (Ping Tox.Message) where
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (Query Tox.Message (FindNode Tox.Message IPv4)) where
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (Query Tox.Message (Ping Tox.Message)) where -- TODO
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (Response Tox.Message (NodeFound Tox.Message IPv4)) where
+    get = error "TODO get"
+    put = error "TODO put"
+instance Serialize (Response Tox.Message (Ping Tox.Message)) where -- TODO
+    get = error "TODO get"
+    put = error "TODO put"
+instance KRPC (Query    Tox.Message (FindNode Tox.Message IPv4))
+              (Response Tox.Message (NodeFound Tox.Message IPv4)) where
+instance KRPC (Query    Tox.Message (Ping Tox.Message ))
+              (Response Tox.Message (Ping Tox.Message )) where
+instance DataHandlers ByteString Tox.Message where
+
+
+-- instance Generic (Response Tox.Message (NodeFound Tox.Message IPv4)) where -- TODO
+
+instance Default Bool where def = False
+
+clientSession :: Node BValue KMessageOf () IPv4 -> Node B.ByteString Tox.Message Bool IPv4 -> MVar () -> Bool ->  RestrictedSocket -> Int -> Handle -> IO ()
+clientSession bt tox signalQuit isBt sock n h = do
     line <- map toLower . dropWhile isSpace <$> hGetLine h
-    let cmd0 action = action >> clientSession st signalQuit sock n h
-        cmd  action = cmd0 $ join $ runDHT st action
+    let dht :: Either (Node BValue KMessageOf () IPv4)
+                      (Node B.ByteString Tox.Message Bool IPv4)
+        dht | isBt      = Left  bt
+            | otherwise = Right tox
+        cmd0 :: IO () -> IO ()
+        cmd0 action = action >> clientSession bt tox signalQuit isBt sock n h
+        cmd :: GenericDHT IPv4 (IO ()) -> IO ()
+        cmd  (GenericDHT action) = cmd0 $ join $ either (flip runDHT action) (flip runDHT action) dht
+        cmd  (BtDHT action)      = cmd0 $ join $ runDHT bt action
         (c,args) = second (dropWhile isSpace) $ break isSpace line
+        switchNetwork dest = clientSession bt tox signalQuit dest sock n h
     case (c,args) of
 
+        ("bt", _) -> switchNetwork True
+
+        ("tox", _) -> switchNetwork False
+
         ("quit", _) -> hPutClient h "" >> hClose h
 
         ("stop", _) -> do hPutClient h "Terminating DHT Daemon."
                           hClose h
                           putMVar signalQuit ()
 
-        ("ls", _) -> cmd $ do
+        ("pid", _) -> cmd0 $ do
+                    pid <- getProcessID
+                    hPutClient h (show pid)
+
+        -- DHT specific
+        ("ls", _) -> cmd $ GenericDHT $ do
                     tbl <- getTable
                     t <- showTable
                     me <- myNodeIdAccordingTo (read "8.8.8.8:6881")
@@ -163,14 +302,15 @@ clientSession st signalQuit sock n h = do
                         , showReport
                             [ ("node-id", show $ pPrint me)
                             , ("internet address", show ip)
-                            , ("buckets", show $ R.shape tbl)]
+                            , ("buckets", show $ R.shape tbl)
+                            , ("network", if isBt then "mainline" else "tox") ]
                         ]
-        ("external-ip", _) -> cmd $ do
+        ("external-ip", _) -> cmd $ BtDHT $ do
                     ip <- routableAddress
                     return $ do
                     hPutClient h $ maybe "" (takeWhile (/=':') . show) ip
 
-        ("swarms", s) -> cmd $ do
+        ("swarms", s) -> cmd $ BtDHT $ do
                     let fltr = case s of
                                     ('-':'v':cs) | all isSpace (take 1 cs)
                                       -> const True
@@ -181,23 +321,21 @@ clientSession st signalQuit sock n h = do
                     return $ do
                     hPutClient h $ showReport r
 
-        ("peers", s) -> cmd $ case readEither s of
+        -- bittorrent only
+        ("peers", s) -> cmd $ BtDHT $ case readEither s of
                         Right ih -> do
                             ps <- allPeers ih
                             seq ih $ return $ do
                             hPutClient h $ showReport $ map (((,) "") . show . pPrint) ps
                         Left er -> return $ hPutClient h er
 
-        ("pid", _) -> cmd $ return $ do
-                    pid <- getProcessID
-                    hPutClient h (show pid)
 #ifdef THREAD_DEBUG
-        ("threads", _) -> cmd $ return $ do
+        ("threads", _) -> cmd0 $ do
                     ts <- threadsInformation
                     tm <- getCurrentTime
                     let r = map (\PerThread{..} -> (show lbl,show (diffUTCTime tm startTime))) ts
                     hPutClient h $ showReport r
-        ("mem", s) -> cmd $ return $ do
+        ("mem", s) -> cmd0 $ do
                     case s of
                         "gc" -> do hPutClient h "Performing garbage collection..."
                                    performMajorGC
@@ -230,24 +368,27 @@ clientSession st signalQuit sock n h = do
                         _ -> hPutClient h "error."
 
 #endif
-        ("closest", s) -> cmd $ do
+        -- DHT specific
+        ("closest", s) -> cmd $ GenericDHT $ do
                         let (ns,hs) = second (dropWhile isSpace) $ break isSpace s
                             parse | null hs = do
                                        ih <- readEither ns
-                                       return (8 :: Int, ih :: InfoHash)
+                                       return (8 :: Int, ih)
                                   | otherwise = do
                                        n <- readEither ns
                                        ih <- readEither hs
-                                       return (n :: Int, ih :: InfoHash)
+                                       return (n :: Int, ih)
                         case parse of
                             Right (n,ih) -> do
+                                nodeIdType ih
                                 tbl <- getTable
                                 let nodes = R.kclosest n ih tbl
                                 return $ do
-                                hPutClient h $ unlines $ map (showEnry . (flip (,) 0)) nodes
+                                hPutClient h $ unlines $ map (showEnry . (flip (,) (error "showEnry"))) nodes
                             Left er -> return $ hPutClient h er
 
-        ("ping", s) -> cmd $ do
+        -- DHT specific
+        ("ping", s) -> cmd $ GenericDHT $ do
                     case readEither s of
                         Right addr -> do result <- try $ pingQ addr
                                          let rs = either (pure . showQueryFail) reportPong result
@@ -255,22 +396,30 @@ clientSession st signalQuit sock n h = do
                                          hPutClient h $ unlines rs
                         Left er -> return $ hPutClient h er
 
-        ("find-nodes", s) -> cmd $ do
+        -- DHT specific
+        ("find-nodes", s) -> cmd $ GenericDHT $ do
                     let (hs,as) = second (dropWhile isSpace) $ break isSpace s
                         parse = do ih <- readEither hs
                                    a <- readEither as
                                    -- XXX: using 'InfoHash' only because 'NodeId' currently
                                    -- has no 'Read' instance.
-                                   return (ih :: InfoHash, a :: NodeAddr IPv4)
+                                   return (ih, a :: NodeAddr IPv4)
                     case parse of
                         Right (ih,a) -> do
-                            result <- try $ queryNode' (a ::NodeAddr IPv4) $ FindNode (R.toNodeId ih)
+                            nodeIdType ih
+                            nodeAddrType a
+                            proxy <- dhtType
+                            let fn = findNodeMessage proxy ih
+                            ipType fn
+                            result <- try $ queryNode' a fn
+                            either (const $ return ()) (\(nid,nf,_) -> nodeIdType nid >> ipType nf) result
                             let rs = either (pure . showQueryFail) reportNodes result
                             return $ do
                             hPutClient h $ unlines rs
                         Left er -> return $ hPutClient h er
 
-        ("get-peers", s) -> cmd $ do
+        -- bittorrent only
+        ("get-peers", s) -> cmd $ BtDHT $ do
                     let (hs,as) = second (dropWhile isSpace) $ break isSpace s
                         parse = do ih <- readEither hs
                                    a <- readEither as
@@ -283,7 +432,8 @@ clientSession st signalQuit sock n h = do
                             hPutClient h $ showReport rs
                         Left er -> return $ hPutClient h er
 
-        ("search-peers", s) -> cmd $ do
+        -- bittorrent only
+        ("search-peers", s) -> cmd $ BtDHT $ do
                     case readEither s of
                         Right ih -> do
                             (tid, s) <- isearch ioGetPeers ih
@@ -315,10 +465,14 @@ consip' (ReflectedIP ip) xs = ("to", show ip) : xs
 
 reportPong (info,myip) = maybe id consip myip [show $ pPrint info]
 
-reportNodes :: (NodeId _, NodeFound IPv4, Maybe ReflectedIP) -> [String]
-reportNodes (nid,NodeFound ns,myip) = maybe id consip myip $ show (pPrint nid) : map (show . pPrint) ns
+reportNodes ::
+    ( Kademlia dht
+    , Pretty (NodeInfo dht ip ())
+    , Pretty (NodeId dht)
+    ) => (NodeId dht, NodeFound dht ip, Maybe ReflectedIP) -> [String]
+reportNodes (nid,ns,myip) = maybe id consip myip $ show (pPrint nid) : map (show . pPrint) (foundNodes ns)
 
-reportPeers :: (NodeId _, GotPeers IPv4, Maybe ReflectedIP) -> [(String,String)]
+reportPeers :: (NodeId KMessageOf, GotPeers IPv4, Maybe ReflectedIP) -> [(String,String)]
 reportPeers (nid,GotPeers r tok,myip)
     = maybe id consip' myip $ ("from", show (pPrint nid))
      : ("token", show tok)
@@ -368,7 +522,7 @@ main = do
     st <- ask
     waitForSignal <- liftIO $ do
         signalQuit <- newEmptyMVar
-        srv <- streamServer (withSession $ clientSession st signalQuit) (SockAddrUnix "dht.sock")
+        srv <- streamServer (withSession $ clientSession st (error "todo: tox state") signalQuit True) (SockAddrUnix "dht.sock")
         return $ liftIO $ do
             () <- takeMVar signalQuit
             quitListening srv
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
diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index 29d4231d..aefd7742 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -88,6 +88,10 @@ module Network.DHT.Mainline
        , Announce  (..)
        , Announced (..)
 #endif
+       , DHTData(..)
+       , SessionTokens(..)
+       , grantToken
+       , checkToken
        ) where
 
 import Control.Applicative
@@ -118,13 +122,19 @@ import Network.DatagramServer.Mainline
 import Data.Maybe
 
 import Data.Torrent (InfoHash)
-import Network.BitTorrent.DHT.Token
+import Network.BitTorrent.DHT.Token as T
+import Network.BitTorrent.DHT.ContactInfo
 #ifdef VERSION_bencoding
 import Network.DatagramServer ()
 #endif
 import Network.DatagramServer.Types hiding (Query,Response)
 import Network.DHT.Types
 import Network.DHT.Routing
+import Data.Time
+import Control.Concurrent.STM
+import System.Random
+import Data.Hashable
+
 
 {-----------------------------------------------------------------------
 -- envelopes
@@ -472,6 +482,59 @@ bep42 addr (NodeId r)
         where msk | BS.length ip == 4  =  ip4mask
                   | otherwise          =  ip6mask
 
+{-----------------------------------------------------------------------
+-- Tokens policy
+-----------------------------------------------------------------------}
+
+data SessionTokens = SessionTokens
+  { tokenMap    :: !TokenMap
+  , lastUpdate  :: !UTCTime
+  , maxInterval :: !NominalDiffTime
+  }
+
+nullSessionTokens :: IO SessionTokens
+nullSessionTokens = SessionTokens
+  <$> (tokens <$> randomIO)
+  <*> 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
+
+{-----------------------------------------------------------------------
+-- Tokens
+-----------------------------------------------------------------------}
+
+tryUpdateSecret :: TVar SessionTokens -> IO ()
+tryUpdateSecret toks = do
+  curTime <- getCurrentTime
+  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)
+
+
+--------------------------
+
+
 instance Kademlia KMessageOf where
     data Ping KMessageOf = Ping
       deriving (Show, Eq, Typeable)
@@ -479,10 +542,17 @@ instance Kademlia KMessageOf where
       deriving (Show, Eq, Typeable)
     newtype NodeFound KMessageOf ip = NodeFound [NodeInfo KMessageOf ip ()]
       deriving (Show, Eq, Typeable)
+    data DHTData KMessageOf ip = TorrentData
+            { contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
+            , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
+            }
+
     pingMessage _ = Ping
     pongMessage _ = Ping
     findNodeMessage _ k = FindNode (toNodeId k)
+    findWho (FindNode nid) = nid
     foundNodes (NodeFound ns) = ns
+    foundNodesMessage ns = NodeFound ns
 
     dhtAdjustID _ fallback ip0 arrival
             = fromMaybe fallback $ do
@@ -494,3 +564,7 @@ instance Kademlia KMessageOf where
 
     namePing      _ = "ping"
     nameFindNodes _ = "find-nodes"
+
+    initializeDHTData = TorrentData
+                <$> newTVarIO def
+                <*> (newTVarIO =<< nullSessionTokens)
diff --git a/src/Network/DHT/Types.hs b/src/Network/DHT/Types.hs
index 79f9e1d3..31ae5948 100644
--- a/src/Network/DHT/Types.hs
+++ b/src/Network/DHT/Types.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE DeriveGeneric #-}
 module Network.DHT.Types
     ( module Network.DHT.Types
     , TableKey
@@ -12,6 +13,7 @@ import Network.Socket (SockAddr)
 import Network.DatagramServer.Types
 import Network.DHT.Routing
 import Data.Typeable
+import GHC.Generics
 
 data TableParameters msg ip u = TableParameters
     { maxBuckets :: Int
@@ -27,7 +29,7 @@ data Query dht a = Query
   { queringNodeId   :: NodeId dht -- ^ node id of /quering/ node;
   , queryIsReadOnly :: Bool   -- ^ node is read-only as per BEP 43
   , queryParams     :: a      -- ^ query parameters.
-  } deriving (Typeable)
+  } deriving (Typeable,Generic)
 
 deriving instance (Eq (NodeId dht), Eq a ) => Eq (Query dht a)
 deriving instance (Show (NodeId dht), Show a ) => Show (Query dht a)
@@ -37,7 +39,7 @@ deriving instance (Show (NodeId dht), Show a ) => Show (Query dht a)
 data Response dht a = Response
   { queredNodeId :: NodeId dht -- ^ node id of /quered/ node;
   , responseVals :: a      -- ^ query result.
-  } deriving (Typeable)
+  } deriving (Typeable,Generic)
 
 deriving instance (Eq (NodeId dht), Eq a ) => Eq (Response dht a)
 deriving instance (Show (NodeId dht), Show a ) => Show (Response dht a)
@@ -51,6 +53,7 @@ class Kademlia dht where
     -- given its ID.
     data FindNode dht ip
     data NodeFound dht ip
+    data DHTData dht ip
     pingMessage :: Proxy dht -> Ping dht
     pongMessage :: Proxy dht -> Ping dht
     findNodeMessage :: TableKey dht k => Proxy dht -> k -> FindNode dht ip
@@ -60,3 +63,4 @@ class Kademlia dht where
     dhtAdjustID :: Address ip => Proxy dht -> NodeId dht -> SockAddr -> Event dht ip u -> NodeId dht
     namePing :: Proxy dht -> QueryMethod dht
     nameFindNodes :: Proxy dht -> QueryMethod dht
+    initializeDHTData :: IO (DHTData dht ip)
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
index 3a2bd020..d0eb136a 100644
--- a/src/Network/KRPC/Method.hs
+++ b/src/Network/KRPC/Method.hs
@@ -89,11 +89,11 @@ showsMethod (Method name) =
 --   @
 --
 class ( Typeable req, Typeable resp
-#ifdef VERSION_bencoding
-      , BEncode req, BEncode resp
-#else
-      , Serialize req, Serialize resp
-#endif
+-- #ifdef VERSION_bencoding
+      -- , BEncode req, BEncode resp
+-- #else
+      -- , Serialize req, Serialize resp
+-- #endif
       )
     => KRPC req resp | req -> resp, resp -> req where