1 files changed, 44 insertions, 45 deletions

diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index d118ceb0..29d4231d 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -123,6 +123,8 @@ import Network.BitTorrent.DHT.Token
 import Network.DatagramServer ()
 #endif
 import Network.DatagramServer.Types hiding (Query,Response)
+import Network.DHT.Types
+import Network.DHT.Routing
 
 {-----------------------------------------------------------------------
 -- envelopes
@@ -140,15 +142,7 @@ read_only_key = "ro"
 
 
 #ifdef VERSION_bencoding
--- | All queries have an \"id\" key and value containing the node ID
--- of the querying node.
-data Query a = Query
-  { queringNodeId   :: NodeId KMessageOf -- ^ node id of /quering/ node;
-  , queryIsReadOnly :: Bool   -- ^ node is read-only as per BEP 43
-  , queryParams     :: a      -- ^ query parameters.
-  } deriving (Show, Eq, Typeable)
-
-instance BEncode a => BEncode (Query a) where
+instance BEncode a => BEncode (Query KMessageOf a) where
   toBEncode Query {..} = toDict $
       BDict.union ( node_id_key .=! queringNodeId
                    .: read_only_key .=? bool Nothing (Just (1 :: Integer)) queryIsReadOnly
@@ -167,14 +161,7 @@ data Query a = Query a
 #endif
 
 #ifdef VERSION_bencoding
--- | All responses have an \"id\" key and value containing the node ID
--- of the responding node.
-data Response a = Response
-  { queredNodeId :: NodeId KMessageOf -- ^ node id of /quered/ node;
-  , responseVals :: a      -- ^ query result.
-  } deriving (Show, Eq, Typeable)
-
-instance BEncode a => BEncode (Response a) where
+instance BEncode a => BEncode (Response KMessageOf a) where
   toBEncode = toBEncode . toQuery
     where
       toQuery (Response nid a) = Query nid False a
@@ -183,28 +170,23 @@ instance BEncode a => BEncode (Response a) where
     where
       fromQuery (Query nid _ a) = Response nid a
 #else
-data Response a = Response a
+data Response KMessageOf a = Response KMessageOf a
 #endif
 
 {-----------------------------------------------------------------------
 -- ping method
 -----------------------------------------------------------------------}
 
--- | The most basic query is a ping. Ping query is used to check if a
+-- / The most basic query is a ping. Ping query is used to check if a
 -- quered node is still alive.
-#ifdef VERSION_bencoding
-data Ping = Ping
-#else
-data Ping = Ping Tox.Nonce8
-#endif
-  deriving (Show, Eq, Typeable)
+-- data Ping = Ping Tox.Nonce8 deriving (Show, Eq, Typeable)
 
 #ifdef VERSION_bencoding
-instance BEncode Ping where
+instance BEncode (Ping KMessageOf) where
   toBEncode Ping = toDict endDict
   fromBEncode _  = pure Ping
 #else
-instance Serialize (Query Ping) where
+instance Serialize (Query (Ping KMessageOf)) where
     get = do
         b <- get
         when ( (b::Word8) /= 0) $ fail "Bad ping request"
@@ -225,7 +207,7 @@ instance Serialize (Response Ping) where
 #endif
 
 -- | \"q\" = \"ping\"
-instance KRPC (Query Ping) (Response Ping) where
+instance KRPC (Query KMessageOf (Ping KMessageOf)) (Response KMessageOf (Ping KMessageOf)) where
 #ifdef VERSION_bencoding
   method = "ping"
 #else
@@ -236,24 +218,20 @@ instance KRPC (Query Ping) (Response Ping) where
 -- find_node method
 -----------------------------------------------------------------------}
 
--- | Find node is used to find the contact information for a node
+-- / Find node is used to find the contact information for a node
 -- given its ID.
-#ifdef VERSION_bencoding
-newtype FindNode ip = FindNode (NodeId KMessageOf)
-#else
-data FindNode ip = FindNode (NodeId Tox.Message) Tox.Nonce8 -- Tox: Get Nodes
-#endif
-  deriving (Show, Eq, Typeable)
+-- data FindNode KMessageOf ip = FindNode (NodeId Tox.Message) Tox.Nonce8 -- Tox: Get Nodes
+  -- deriving (Show, Eq, Typeable)
 
 target_key :: BKey
 target_key = "target"
 
 #ifdef VERSION_bencoding
-instance Typeable ip => BEncode (FindNode ip) where
+instance Typeable ip => BEncode (FindNode KMessageOf ip) where
   toBEncode (FindNode nid) = toDict   $ target_key .=! nid .: endDict
   fromBEncode              = fromDict $ FindNode  <$>! target_key
 #else
-instance Serialize (Query (FindNode ip)) where
+instance Serialize (Query KMessageOf (FindNode KMessageOf ip)) where
     get = do
         nid <- get
         nonce <- get
@@ -268,12 +246,11 @@ instance Serialize (Query (FindNode ip)) where
 -- nodes in its own routing table.
 --
 #ifdef VERSION_bencoding
-newtype NodeFound ip = NodeFound [NodeInfo KMessageOf ip ()]
+-- newtype NodeFound KMessageOf ip = NodeFound [NodeInfo KMessageOf ip ()] deriving (Show, Eq, Typeable)
 #else
-data NodeFound ip = NodeFound [Tox.NodeFormat] Tox.Nonce8
+data NodeFound KMessageOf ip = NodeFound [Tox.NodeFormat] Tox.Nonce8 deriving (Show, Eq, Typeable)
 #endif
 -- Tox: send_nodes
-  deriving (Show, Eq, Typeable)
 
 nodes_key :: BKey
 nodes_key = "nodes"
@@ -290,7 +267,7 @@ binary k = field (req k) >>= either (fail . format) return .
   where
     format str = "fail to deserialize " ++ show k ++ " field: " ++ str
 
-instance Address ip => BEncode (NodeFound ip) where
+instance Address ip => BEncode (NodeFound KMessageOf ip) where
   toBEncode (NodeFound ns) = toDict $
        nodes_key .=! runPut (mapM_ put ns)
     .: endDict
@@ -298,7 +275,7 @@ instance Address ip => BEncode (NodeFound ip) where
   -- TODO: handle IPv6 by reading the "nodes6" key (see bep 32)
   fromBEncode bval = NodeFound <$> (traverse from4 =<< fromDict (binary nodes_key) bval)
 #else
-instance Serialize (Response (NodeFound ip)) where
+instance Serialize (Response KMessageOf (NodeFound KMessageOf ip)) where
     get = do
         count <- get :: Get Word8
         nodes <- sequence $ replicate (fromIntegral count) get
@@ -314,7 +291,7 @@ instance Serialize (Response (NodeFound ip)) where
 
 -- | \"q\" == \"find_node\"
 instance (Address ip, Typeable ip)
-      => KRPC (Query (FindNode ip)) (Response (NodeFound ip)) where
+      => KRPC (Query KMessageOf (FindNode KMessageOf ip)) (Response KMessageOf (NodeFound KMessageOf ip)) where
 #ifdef VERSION_bencoding
   method = "find_node"
 #else
@@ -383,7 +360,7 @@ instance (Typeable ip, Serialize ip) => BEncode (GotPeers ip) where
 
 -- | \"q" = \"get_peers\"
 instance (Typeable ip, Serialize ip) =>
-         KRPC (Query (GetPeers ip)) (Response (GotPeers ip)) where
+         KRPC (Query KMessageOf (GetPeers ip)) (Response KMessageOf (GotPeers ip)) where
   method = "get_peers"
 
 {-----------------------------------------------------------------------
@@ -455,7 +432,7 @@ instance BEncode Announced where
   fromBEncode _ = pure  Announced
 
 -- | \"q" = \"announce\"
-instance KRPC (Query Announce) (Response Announced) where
+instance KRPC (Query KMessageOf Announce) (Response KMessageOf Announced) where
   method = "announce_peer"
 
 -- endif VERSION_bencoding
@@ -495,3 +472,25 @@ bep42 addr (NodeId r)
         where msk | BS.length ip == 4  =  ip4mask
                   | otherwise          =  ip6mask
 
+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)
+    pingMessage _ = Ping
+    pongMessage _ = Ping
+    findNodeMessage _ k = FindNode (toNodeId k)
+    foundNodes (NodeFound ns) = ns
+
+    dhtAdjustID _ fallback ip0 arrival
+            = fromMaybe fallback $ do
+                    ip <- fromSockAddr ip0 -- :: Maybe ip
+                    let _ = ip `asTypeOf` nodeAddr (foreignNode arrival)
+                    listToMaybe
+                        $ rank id (nodeId $ foreignNode arrival)
+                        $ bep42s ip fallback
+
+    namePing      _ = "ping"
+    nameFindNodes _ = "find-nodes"