1 files changed, 11 insertions, 9 deletions

diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 0e2bfdd9..c3df683a 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -93,15 +93,14 @@ import Data.Bool
 #ifdef VERSION_bencoding
 import Data.BEncode as BE
 import Data.BEncode.BDict as BDict
-import Network.BitTorrent.Address
 #else
 import qualified Data.Tox as Tox
 import Data.Tox (NodeId)
 import Data.Word
 import Control.Monad
 import Network.KRPC.Method
-import Network.BitTorrent.Address hiding (NodeId)
 #endif
+import Network.BitTorrent.Address hiding (NodeId)
 import Data.ByteString (ByteString)
 import Data.List as L
 import Data.Monoid
@@ -109,11 +108,14 @@ import Data.Serialize as S
 import Data.Typeable
 import Network
 import Network.KRPC
+import Network.KRPC.Message (KMessageOf)
 import Data.Maybe
 
 import Data.Torrent (InfoHash)
 import Network.BitTorrent.DHT.Token
 import Network.KRPC ()
+import Network.DHT.Mainline ()
+import Network.RPC hiding (Query,Response)
 
 {-----------------------------------------------------------------------
 -- envelopes
@@ -134,7 +136,7 @@ read_only_key = "ro"
 -- | All queries have an \"id\" key and value containing the node ID
 -- of the querying node.
 data Query a = Query
-  { queringNodeId   :: NodeId -- ^ node id of /quering/ node;
+  { 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)
@@ -161,7 +163,7 @@ data Query a = Query a
 -- | All responses have an \"id\" key and value containing the node ID
 -- of the responding node.
 data Response a = Response
-  { queredNodeId :: NodeId -- ^ node id of /quered/ node;
+  { queredNodeId :: NodeId KMessageOf -- ^ node id of /quered/ node;
   , responseVals :: a      -- ^ query result.
   } deriving (Show, Eq, Typeable)
 
@@ -233,7 +235,7 @@ instance KRPC (Query Ping) (Response Ping) where
 -- | Find node is used to find the contact information for a node
 -- given its ID.
 #ifdef VERSION_bencoding
-newtype FindNode ip = FindNode NodeId
+newtype FindNode ip = FindNode (NodeId KMessageOf)
 #else
 data FindNode ip = FindNode NodeId Tox.Nonce8 -- Tox: Get Nodes
 #endif
@@ -262,7 +264,7 @@ instance Serialize (Query (FindNode ip)) where
 -- nodes in its own routing table.
 --
 #ifdef VERSION_bencoding
-newtype NodeFound ip = NodeFound [NodeInfo ip]
+newtype NodeFound ip = NodeFound [NodeInfo KMessageOf ip ()]
 #else
 data NodeFound ip = NodeFound [Tox.NodeFormat] Tox.Nonce8
 #endif
@@ -273,9 +275,9 @@ nodes_key :: BKey
 nodes_key = "nodes"
 
 -- Convert IPv4 address.  Useful for using variadic IP type.
-from4 :: forall s. Address s => NodeInfo IPv4 -> Either String (NodeInfo s)
+from4 :: forall dht u s. Address s => NodeInfo dht IPv4 u -> Either String (NodeInfo dht s u)
 from4 n = maybe (Left "Error converting IPv4") Right
-                    $ traverse (fromAddr :: IPv4 -> Maybe s) n
+                    $ traverseAddress (fromAddr :: IPv4 -> Maybe s) n
 
 #ifdef VERSION_bencoding
 binary :: Serialize a => BKey -> BE.Get [a]
@@ -334,7 +336,7 @@ instance Typeable ip => BEncode (GetPeers ip) where
   toBEncode (GetPeers ih) = toDict   $ info_hash_key .=! ih .: endDict
   fromBEncode             = fromDict $ GetPeers <$>! info_hash_key
 
-type PeerList ip = Either [NodeInfo ip] [PeerAddr ip]
+type PeerList ip = Either [NodeInfo KMessageOf ip ()] [PeerAddr ip]
 
 data GotPeers ip = GotPeers
   { -- | If the queried node has no peers for the infohash, returned