3 files changed, 31 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 44dc9b2f..0e2bfdd9 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -61,6 +61,7 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE UndecidableInstances  #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TypeFamilies #-}
 module Network.BitTorrent.DHT.Message
        ( -- * Envelopes
          Query (..)
@@ -217,6 +218,9 @@ instance Serialize (Response Ping) where
 -- | \"q\" = \"ping\"
 instance KRPC (Query Ping) (Response Ping) where
 #ifdef VERSION_bencoding
+  type Envelope (Query Ping) (Response Ping) = BValue
+  seal = toBEncode
+  unseal = fromBEncode
   method = "ping"
 #else
   method = Method Tox.Ping -- response: Tox.Pong
@@ -229,9 +233,9 @@ 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 = FindNode NodeId
+newtype FindNode ip = FindNode NodeId
 #else
-data FindNode = FindNode NodeId Tox.Nonce8 -- Tox: Get Nodes
+data FindNode ip = FindNode NodeId Tox.Nonce8 -- Tox: Get Nodes
 #endif
   deriving (Show, Eq, Typeable)
 
@@ -239,11 +243,11 @@ target_key :: BKey
 target_key = "target"
 
 #ifdef VERSION_bencoding
-instance BEncode FindNode where
+instance Typeable ip => BEncode (FindNode ip) where
   toBEncode (FindNode nid) = toDict   $ target_key .=! nid .: endDict
   fromBEncode              = fromDict $ FindNode  <$>! target_key
 #else
-instance Serialize (Query FindNode) where
+instance Serialize (Query (FindNode ip)) where
     get = do
         nid <- get
         nonce <- get
@@ -304,8 +308,11 @@ instance Serialize (Response (NodeFound ip)) where
 
 -- | \"q\" == \"find_node\"
 instance (Address ip, Typeable ip)
-      => KRPC (Query FindNode) (Response (NodeFound ip)) where
+      => KRPC (Query (FindNode ip)) (Response (NodeFound ip)) where
 #ifdef VERSION_bencoding
+  type Envelope (Query (FindNode ip)) (Response (NodeFound ip)) = BValue
+  seal = toBEncode
+  unseal = fromBEncode
   method = "find_node"
 #else
   method = Method Tox.GetNodes -- response: Tox.SendNodes
@@ -317,13 +324,13 @@ instance (Address ip, Typeable ip)
 -----------------------------------------------------------------------}
 
 -- | Get peers associated with a torrent infohash.
-newtype GetPeers = GetPeers InfoHash
+newtype GetPeers ip = GetPeers InfoHash
   deriving (Show, Eq, Typeable)
 
 info_hash_key :: BKey
 info_hash_key = "info_hash"
 
-instance BEncode GetPeers where
+instance Typeable ip => BEncode (GetPeers ip) where
   toBEncode (GetPeers ih) = toDict   $ info_hash_key .=! ih .: endDict
   fromBEncode             = fromDict $ GetPeers <$>! info_hash_key
 
@@ -373,7 +380,10 @@ instance (Typeable ip, Serialize ip) => BEncode (GotPeers ip) where
 
 -- | \"q" = \"get_peers\"
 instance (Typeable ip, Serialize ip) =>
-         KRPC (Query GetPeers) (Response (GotPeers ip)) where
+         KRPC (Query (GetPeers ip)) (Response (GotPeers ip)) where
+  type Envelope (Query (GetPeers ip)) (Response (GotPeers ip)) = BValue
+  seal = toBEncode
+  unseal = fromBEncode
   method = "get_peers"
 
 {-----------------------------------------------------------------------
@@ -446,6 +456,9 @@ instance BEncode Announced where
 
 -- | \"q" = \"announce\"
 instance KRPC (Query Announce) (Response Announced) where
+  type Envelope (Query Announce) (Response Announced) = BValue
+  seal = toBEncode
+  unseal = fromBEncode
   method = "announce_peer"
 
 -- endif VERSION_bencoding
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index c7e48920..a1934014 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -14,6 +14,7 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell     #-}
 {-# LANGUAGE TupleSections       #-}
+{-# LANGUAGE GADTs #-}
 module Network.BitTorrent.DHT.Query
        ( -- * Handler
          -- | To bind specific set of handlers you need to pass
@@ -71,6 +72,7 @@ import Data.Either
 import Data.List as L
 import Data.Monoid
 import Data.Text as T
+import Data.BEncode (BValue)
 import qualified Data.Set as Set
          ;import Data.Set (Set)
 import Network
@@ -89,14 +91,17 @@ import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 import Control.Concurrent.STM
 import qualified Network.BitTorrent.DHT.Search as Search
+import Network.DHT.Mainline
 
 {-----------------------------------------------------------------------
 --  Handlers
 -----------------------------------------------------------------------}
 
-nodeHandler :: Address ip => KRPC (Query a) (Response b)
+nodeHandler :: ( Address ip
+               , KRPC (Query a) (Response b)
+               , Envelope (Query a) (Response b) ~ BValue )
            => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
-nodeHandler action = handler $ \ sockAddr qry -> do
+nodeHandler action = handler mainline $ \ sockAddr qry -> do
   let remoteId  = queringNodeId qry
       read_only = queryIsReadOnly qry
       q         = queryParams qry
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 82926b28..db8e7cff 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -106,6 +106,8 @@ import Data.Serialize as S
 import Data.Torrent as Torrent
 import Network.KRPC as KRPC hiding (Options, def)
 import qualified Network.KRPC as KRPC (def)
+import Network.KRPC.Message (KMessageOf)
+import Data.BEncode (BValue)
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.ContactInfo (PeerStore)
 import qualified Network.BitTorrent.DHT.ContactInfo as P
@@ -312,7 +314,7 @@ instance MonadLogger (DHT ip) where
     logger <- asks loggerFun
     liftIO $ logger loc src lvl (toLogStr msg)
 
-type NodeHandler ip = Handler (DHT ip)
+type NodeHandler ip = Handler (DHT ip) KMessageOf BValue
 
 -- | Run DHT session. You /must/ properly close session using
 -- 'closeNode' function, otherwise socket or other scarce resources may