1 files changed, 2 insertions, 2 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index a7359bda..c0413322 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -444,7 +444,7 @@ encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr)
 encodePacket msg ni = ( toStrict $ BE.encode msg
                       , nodeAddr ni )
 
-classify :: Message BValue -> MessageClass String Method TransactionId
+classify :: Message BValue -> MessageClass String Method TransactionId NodeInfo (Message BValue)
 classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid
 classify (R { msgID = tid                   }) = IsResponse tid
 
@@ -566,7 +566,7 @@ newClient swarms addr = do
         outgoingClient = client { clientNet = net { awaitMessage = ($ Nothing) } }
 
         dispatch = DispatchMethods
-            { classifyInbound = classify -- :: x -> MessageClass err meth tid
+            { classifyInbound = classify -- :: x -> MessageClass err meth tid addr x
             , lookupHandler   = handlers -- :: meth -> Maybe (MethodHandler err tid addr x)
             , tableMethods    = mapT     -- :: TransactionMethods tbl tid x
             }