3 files changed, 54 insertions, 43 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 68f91446..e4716d1a 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -10,6 +10,7 @@
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TypeFamilies #-}
 
 import Control.Arrow
 import Control.Monad
@@ -223,6 +224,8 @@ ipType :: f dht ip -> DHT raw dht u ip ()
 ipType _ = return ()
 
 instance Kademlia Tox.Message where
+    data DHTData Tox.Message ip = ToxData
+    initializeDHTData = return ToxData
 instance Pretty (NodeId Tox.Message) where
 instance Pretty (NodeInfo Tox.Message IPv4 ()) where
 instance Pretty (NodeInfo Tox.Message IPv4 Bool) where -- TODO
@@ -230,32 +233,32 @@ instance Read (NodeId KMessageOf) where
     readsPrec d s = map (\(ih,s) -> (toNodeId (ih::InfoHash),s)) $ readsPrec d s
 instance Read (NodeId Tox.Message) where -- TODO
 instance Serialize (FindNode Tox.Message IPv4) where
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 1"
+    put = error "TODO put 2"
 instance Serialize (NodeFound Tox.Message IPv4) where
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 3"
+    put = error "TODO put 4"
 instance Serialize (Ping Tox.Message) where
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 5"
+    put = error "TODO put 6"
 instance Serialize (Query Tox.Message (FindNode Tox.Message IPv4)) where
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 7"
+    put = error "TODO put 8"
 instance Serialize (Query Tox.Message (Ping Tox.Message)) where -- TODO
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 9"
+    put = error "TODO put 10"
 instance Serialize (Response Tox.Message (NodeFound Tox.Message IPv4)) where
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 11"
+    put = error "TODO put 12"
 instance Serialize (Response Tox.Message (Ping Tox.Message)) where -- TODO
-    get = error "TODO get"
-    put = error "TODO put"
+    get = error "TODO get 13"
+    put = error "TODO put 14"
 instance KRPC Tox.Message (Query    Tox.Message (FindNode Tox.Message IPv4))
               (Response Tox.Message (NodeFound Tox.Message IPv4)) where
-    method = error "TODO method"
+    method = error "TODO method 15"
 instance KRPC Tox.Message (Query    Tox.Message (Ping Tox.Message ))
               (Response Tox.Message (Ping Tox.Message )) where
-    method = error "TODO method"
+    method = error "TODO method 16"
 instance DataHandlers ByteString Tox.Message where
 
 
@@ -490,6 +493,9 @@ main = do
                 ["-p",port] | not ("-" `isPrefixOf` port) -> return port
                 ("-p":_)      -> error "Port not specified! (-p PORT)"
                 _             -> defaultPort
+
+    tox_state <- godht (show (succ (read p::Int))) $ \a me0 -> ask
+
     godht p $ \a me0 -> do
     printTable
     bs <- liftIO bootstrapNodes
@@ -525,7 +531,7 @@ main = do
     st <- ask
     waitForSignal <- liftIO $ do
         signalQuit <- newEmptyMVar
-        srv <- streamServer (withSession $ clientSession st (error "todo: tox state") signalQuit True) (SockAddrUnix "dht.sock")
+        srv <- streamServer (withSession $ clientSession st tox_state signalQuit True) (SockAddrUnix "dht.sock")
         return $ liftIO $ do
             () <- takeMVar signalQuit
             quitListening srv
diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index b756ff6a..6ef6d450 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -438,7 +438,7 @@ data Announced = Announced
   deriving (Show, Eq, Typeable)
 
 instance BEncode Announced where
-  toBEncode   _ = toBEncode Ping
+  toBEncode   _ = toBEncode ( Ping :: Ping KMessageOf )
   fromBEncode _ = pure  Announced
 
 -- | \"q" = \"announce\"
@@ -536,23 +536,11 @@ checkToken sessionTokens addr questionableToken = do
 
 
 instance Kademlia KMessageOf where
-    data Ping KMessageOf = Ping
-      deriving (Show, Eq, Typeable)
-    newtype FindNode KMessageOf ip = FindNode (NodeId KMessageOf)
-      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
diff --git a/src/Network/DHT/Types.hs b/src/Network/DHT/Types.hs
index 31ae5948..91d1fc22 100644
--- a/src/Network/DHT/Types.hs
+++ b/src/Network/DHT/Types.hs
@@ -44,22 +44,39 @@ data Response dht a = Response
 deriving instance (Eq (NodeId dht), Eq a ) => Eq (Response dht a)
 deriving instance (Show (NodeId dht), Show a ) => Show (Response dht a)
 
+-- | The most basic query is a ping. Ping query is used to check if a
+-- quered node is still alive.
+data Ping ( dht :: * -> *) = Ping
+  deriving (Show, Eq, Typeable)
+-- | Find node is used to find the contact information for a node
+-- given its ID.
+newtype FindNode dht ip = FindNode (NodeId dht)
+  deriving (Typeable)
+newtype NodeFound dht ip = NodeFound [NodeInfo dht ip ()]
+  deriving (Typeable)
+
+deriving instance Eq (NodeId dht) => Eq (FindNode dht ip)
+deriving instance Eq (NodeId dht) => Eq (NodeFound dht ip)
+deriving instance Show (NodeId dht) => Show (FindNode dht ip)
+deriving instance ( Show (NodeId dht)
+                  , Show ip
+                  ) => Show (NodeFound dht ip)
+
+pingMessage :: Proxy dht -> Ping dht
+pingMessage _ = Ping
+pongMessage :: Proxy dht -> Ping dht
+pongMessage _ = Ping
+findNodeMessage :: TableKey dht k => Proxy dht -> k -> FindNode dht ip
+findNodeMessage _ k = FindNode (toNodeId k)
+foundNodesMessage :: [NodeInfo dht ip ()] -> NodeFound dht ip
+findWho (FindNode nid) = nid
+foundNodes :: NodeFound dht ip -> [NodeInfo dht ip ()]
+foundNodes (NodeFound ns) = ns
+findWho :: FindNode dht ip -> NodeId dht
+foundNodesMessage ns = NodeFound ns
 
 class Kademlia dht where
-    -- | The most basic query is a ping. Ping query is used to check if a
-    -- quered node is still alive.
-    data Ping dht
-    -- | Find node is used to find the contact information for a node
-    -- 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
-    foundNodesMessage :: [NodeInfo dht ip ()] -> NodeFound dht ip
-    foundNodes :: NodeFound dht ip -> [NodeInfo dht ip ()]
-    findWho :: FindNode dht ip -> NodeId dht
     dhtAdjustID :: Address ip => Proxy dht -> NodeId dht -> SockAddr -> Event dht ip u -> NodeId dht
     namePing :: Proxy dht -> QueryMethod dht
     nameFindNodes :: Proxy dht -> QueryMethod dht