2 files changed, 25 insertions, 2 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index a1934014..e1104cb9 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -72,7 +72,6 @@ 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
@@ -85,13 +84,20 @@ import Network.KRPC hiding (Options, def)
 import Network.KRPC.Message (ReflectedIP(..))
 import Network.KRPC.Manager (QueryFailure(..))
 import Data.Torrent
-import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 import Control.Concurrent.STM
 import qualified Network.BitTorrent.DHT.Search as Search
+#ifdef VERSION_bencoding
+import Network.BitTorrent.Address
+import Data.BEncode (BValue)
 import Network.DHT.Mainline
+#else
+import Network.BitTorrent.Address hiding (NodeId)
+import Data.ByteString (ByteString)
+import Data.Tox
+#endif
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -99,12 +105,23 @@ import Network.DHT.Mainline
 
 nodeHandler :: ( Address ip
                , KRPC (Query a) (Response b)
+#ifdef VERSION_bencoding
                , Envelope (Query a) (Response b) ~ BValue )
+#else
+               , Envelope (Query a) (Response b) ~ ByteString )
+#endif
            => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
+#ifdef VERSION_bencoding
 nodeHandler action = handler mainline $ \ sockAddr qry -> do
   let remoteId  = queringNodeId qry
       read_only = queryIsReadOnly qry
       q         = queryParams qry
+#else
+nodeHandler action = handler (error "TODO TOX Messaging") $ \ sockAddr qry -> do
+  let remoteId  = msgClient qry
+      read_only = False
+      q         = msgPayload qry
+#endif
   case fromSockAddr sockAddr of
     Nothing    -> throwIO BadAddress
     Just naddr -> do
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index db8e7cff..20dba595 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -107,7 +107,9 @@ 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)
+#ifdef VERSION_bencoding
 import Data.BEncode (BValue)
+#endif
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.ContactInfo (PeerStore)
 import qualified Network.BitTorrent.DHT.ContactInfo as P
@@ -314,7 +316,11 @@ instance MonadLogger (DHT ip) where
     logger <- asks loggerFun
     liftIO $ logger loc src lvl (toLogStr msg)
 
+#ifdef VERSION_bencoding
 type NodeHandler ip = Handler (DHT ip) KMessageOf BValue
+#else
+type NodeHandler ip = Handler (DHT ip) KMessageOf ByteString
+#endif
 
 -- | Run DHT session. You /must/ properly close session using
 -- 'closeNode' function, otherwise socket or other scarce resources may