1 files changed, 33 insertions, 15 deletions

diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index 7cd33c0d..540b74f9 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -1,22 +1,40 @@
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE TypeFamilies #-}
 module Network.DHT.Mainline where
 
+import Network.Socket
 import Network.RPC
-import Network.KRPC.Message
+import Network.KRPC.Message as KRPC
 import Data.BEncode as BE
 import qualified Data.ByteString.Lazy as L
+import Network.BitTorrent.Address as BT (NodeId)
 
-mainline :: Messaging KMessageOf TransactionId BValue
-mainline = Messaging
-    { messageClass = \case Q _ -> Query
-                           R _ -> Response
-                           E _ -> Error
-    , messageTransaction = \case Q q -> queryId q
-                                 R r -> respId r
-                                 E e -> errorId e
-    , messagePayload = \case Q q -> queryArgs q
-                             R r -> respVals r
-                             E e -> error "TODO: messagePayload for KError"
-    , encodePayload = fmap (L.toStrict . BE.encode)
-    , decodePayload = sequence . fmap BE.decode
-    }
+instance Envelope KMessageOf where
+    type TransactionID KMessageOf = KRPC.TransactionId
+    type NodeId KMessageOf = BT.NodeId
+
+    envelopePayload (Q q) = queryArgs q
+    envelopePayload (R r) = respVals r
+    envelopePayload (E _) = error "TODO: messagePayload for KError"
+
+    envelopeTransaction (Q q) = queryId q
+    envelopeTransaction (R r) = respId r
+    envelopeTransaction (E e) = errorId e
+
+    envelopeClass (Q _) = Query
+    envelopeClass (R _) = Response
+    envelopeClass (E _) = Error
+
+    buildReply self addr qry response =
+        (R (KResponse response (envelopeTransaction qry) (Just $ ReflectedIP addr)))
+
+instance WireFormat BValue KMessageOf where
+    type SerializableTo BValue = BEncode
+    type CipherContext BValue KMessageOf = ()
+
+    decodeHeaders _ bs = BE.decode bs >>= BE.fromBEncode
+    decodePayload kmsg = mapM BE.fromBEncode kmsg
+
+    encodeHeaders _ kmsg = L.toStrict $ BE.encode kmsg
+    encodePayload msg = fmap BE.toBEncode msg