8 files changed, 127 insertions, 61 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
diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
new file mode 100644
index 00000000..7cd33c0d
--- /dev/null
+++ b/src/Network/DHT/Mainline.hs
@@ -0,0 +1,22 @@
+{-# LANGUAGE LambdaCase #-}
+module Network.DHT.Mainline where
+
+import Network.RPC
+import Network.KRPC.Message
+import Data.BEncode as BE
+import qualified Data.ByteString.Lazy as L
+
+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
+    }
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index e7f0563b..b1e93101 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -76,6 +76,7 @@ import Data.Text as T
 import Data.Text.Encoding as T
 import Data.Tuple
 import Data.Typeable
+import Network.RPC
 import Network.KRPC.Message
 import Network.KRPC.Method
 import Network.Socket hiding (listen)
@@ -136,11 +137,11 @@ type CallId             = (TransactionId, SockAddr)
 type CallRes            = MVar (KQueryArgs, KResult) -- (raw response, decoded response)
 type PendingCalls       = IORef (Map CallId CallRes)
 
-type HandlerBody h = SockAddr -> KQueryArgs -> h (Either String KQueryArgs)
+type HandlerBody h msg v = SockAddr -> msg v -> h (Either String v)
 
 -- | Handler is a function which will be invoked then some /remote/
 -- node querying /this/ node.
-type Handler     h = (MethodName, HandlerBody h)
+type Handler     h msg v = (MethodName, HandlerBody h msg v)
 
 -- | Keep track pending queries made by /this/ node and handle queries
 -- made by /remote/ nodes.
@@ -150,7 +151,7 @@ data Manager h = Manager
   , listenerThread     :: !(MVar ThreadId)
   , transactionCounter :: {-# UNPACK #-} !TransactionCounter
   , pendingCalls       :: {-# UNPACK #-} !PendingCalls
-  , handlers           :: [Handler h]
+  , handlers           :: [Handler h KMessageOf BValue]
   }
 
 -- | A monad which can perform or handle queries.
@@ -182,10 +183,10 @@ sockAddrFamily (SockAddrCan _        ) = AF_CAN
 
 -- | Bind socket to the specified address. To enable query handling
 -- run 'listen'.
-newManager :: Options        -- ^ various protocol options;
-           -> SockAddr       -- ^ address to listen on;
-           -> [Handler h]    -- ^ handlers to run on incoming queries.
-           -> IO (Manager h) -- ^ new rpc manager.
+newManager :: Options                       -- ^ various protocol options;
+           -> SockAddr                      -- ^ address to listen on;
+           -> [Handler h KMessageOf BValue] -- ^ handlers to run on incoming queries.
+           -> IO (Manager h)                -- ^ new rpc manager.
 newManager opts @ Options {..} servAddr handlers = do
     validateOptions opts
     sock  <- bindServ
@@ -217,7 +218,7 @@ isActive Manager {..} = liftIO $ isBound sock
 
 -- | Normally you should use Control.Monad.Trans.Resource.allocate
 -- function.
-withManager :: Options -> SockAddr -> [Handler h]
+withManager :: Options -> SockAddr -> [Handler h KMessageOf BValue]
             -> (Manager h -> IO a) -> IO a
 withManager opts addr hs = bracket (newManager opts addr hs) closeManager
 
@@ -408,35 +409,25 @@ prettyQF e = T.encodeUtf8 $ "handler fail while performing query: "
 --   If the handler make some 'query' normally it /should/ handle
 --   corresponding 'QueryFailure's.
 --
-handler :: forall h a b. (KRPC a b, Monad h)
-        => (SockAddr -> a -> h b) -> Handler h
-handler body = (name, wrapper)
+handler :: forall h a b msg. (KRPC a b, Applicative h, Functor msg)
+        => Messaging msg TransactionId (Envelope a b) -> (SockAddr -> a -> h b) -> Handler h msg (Envelope a b)
+handler msging body = (name, wrapper)
   where
     Method name = method :: Method a b
     wrapper addr args =
-#ifdef VERSION_bencoding
-      case fromBEncode args of
-#else
-      case S.decode args of
-#endif
-        Left  e -> return $ Left e
-        Right a -> do
-          r <- body addr a
-#ifdef VERSION_bencoding
-          return $ Right $ toBEncode r
-#else
-          return $ Right $ S.encode r
-#endif
+      case unseal $ messagePayload msging args of
+        Left  e -> pure $ Left e
+        Right a -> Right . seal <$> body addr a
 
 runHandler :: MonadKRPC h m
-           => HandlerBody h -> SockAddr -> KQuery -> m KResult
+           => HandlerBody h KMessageOf BValue -> SockAddr -> KQuery -> m KResult
 runHandler h addr m = Lifted.catches wrapper failbacks
   where
     signature = querySignature (queryMethod m) (queryId m) addr
 
     wrapper = do
       $(logDebugS) "handler.quered" signature
-      result <- liftHandler (h addr (queryArgs m))
+      result <- liftHandler (h addr (Q m))
 
       case result of
         Left msg -> do
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index d48fa8ac..19f9fc9e 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -13,14 +13,16 @@
 --   See <http://www.bittorrent.org/beps/bep_0005.html#krpc-protocol>
 --
 {-# LANGUAGE CPP                    #-}
-{-# LANGUAGE OverloadedStrings      #-}
+{-# LANGUAGE DefaultSignatures      #-}
+{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE DeriveFunctor          #-}
+{-# LANGUAGE DeriveTraversable      #-}
 {-# LANGUAGE FlexibleContexts       #-}
 {-# LANGUAGE FlexibleInstances      #-}
-{-# LANGUAGE TypeSynonymInstances   #-}
-{-# LANGUAGE MultiParamTypeClasses  #-}
 {-# LANGUAGE FunctionalDependencies #-}
-{-# LANGUAGE DefaultSignatures      #-}
-{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE MultiParamTypeClasses  #-}
+{-# LANGUAGE OverloadedStrings      #-}
+{-# LANGUAGE TypeSynonymInstances   #-}
 module Network.KRPC.Message
        ( -- * Transaction
          TransactionId
@@ -32,7 +34,8 @@ module Network.KRPC.Message
        , unknownMessage
 
          -- * Query
-       , KQuery(..)
+       , KQueryOf(..)
+       , KQuery
 #ifndef VERSION_bencoding
        , queryArgs
        , queryMethod
@@ -41,11 +44,13 @@ module Network.KRPC.Message
        , MethodName
 
          -- * Response
-       , KResponse(..)
+       , KResponseOf(..)
+       , KResponse
        , ReflectedIP(..)
 
          -- * Message
-       , KMessage (..)
+       , KMessageOf (..)
+       , KMessage
        , KQueryArgs
 
        ) where
@@ -208,11 +213,13 @@ type KQueryArgs = ByteString      -- msgPayload
 -- callee and pass arguments in. Therefore query may be only sent from
 -- client to server but not in the opposite direction.
 --
-data KQuery = KQuery
-  { queryArgs   :: !KQueryArgs     -- ^ values to be passed to method;
+data KQueryOf a = KQuery
+  { queryArgs   :: !a              -- ^ values to be passed to method;
   , queryMethod :: !MethodName     -- ^ method to call;
   , queryId     :: !TransactionId  -- ^ one-time query token.
-  } deriving ( Show, Eq, Ord, Typeable, Read )
+  } deriving ( Show, Eq, Ord, Typeable, Read, Functor, Foldable, Traversable )
+
+type KQuery = KQueryOf KQueryArgs
 
 -- | Queries, or KRPC message dictionaries with a \"y\" value of
 -- \"q\", contain two additional keys; \"q\" and \"a\". Key \"q\" has
@@ -223,7 +230,7 @@ data KQuery = KQuery
 --
 --    > { "t" : "aa", "y" : "q", "q" : "ping", "a" : { "msg" : "hi!" } }
 --
-instance BEncode KQuery where
+instance (Typeable a, BEncode a) => BEncode (KQueryOf a) where
   toBEncode KQuery {..} = toDict $
        "a" .=! queryArgs
     .: "q" .=! queryMethod
@@ -288,11 +295,13 @@ encodeAddr _ = B.empty
 --   * KResponse can be only sent from server to client.
 --
 #ifdef VERSION_bencoding
-data KResponse = KResponse
-  { respVals :: KQueryArgs     -- ^ 'BDict' containing return values;
+data KResponseOf a = KResponse
+  { respVals :: a              -- ^ 'BDict' containing return values;
   , respId   :: TransactionId  -- ^ match to the corresponding 'queryId'.
   , respIP   :: Maybe ReflectedIP
-  } deriving (Show, Eq, Ord, Typeable)
+  } deriving (Show, Eq, Ord, Typeable, Functor, Foldable, Traversable)
+
+type KResponse = KResponseOf KQueryArgs
 
 -- | Responses, or KRPC message dictionaries with a \"y\" value of
 -- \"r\", contain one additional key \"r\". The value of \"r\" is a
@@ -302,7 +311,7 @@ data KResponse = KResponse
 --
 --   > { "t" : "aa", "y" : "r", "r" : { "msg" : "you've sent: hi!" } }
 --
-instance BEncode KResponse where
+instance (Typeable a, BEncode a) => BEncode (KResponseOf a) where
   toBEncode KResponse {..} = toDict $
        "ip" .=? respIP
     .: "r" .=! respVals
@@ -329,11 +338,13 @@ respIP         = Nothing :: Maybe ReflectedIP
 
 #ifdef VERSION_bencoding
 -- | Generic KRPC message.
-data KMessage
-  = Q KQuery
-  | R KResponse
+data KMessageOf a
+  = Q (KQueryOf a)
+  | R (KResponseOf a)
   | E KError
-    deriving (Show, Eq)
+    deriving (Show, Eq, Functor, Foldable, Traversable)
+
+type KMessage = KMessageOf KQueryArgs
 
 instance BEncode KMessage where
   toBEncode (Q q) = toBEncode q
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
index 2a791924..ad93cb8b 100644
--- a/src/Network/KRPC/Method.hs
+++ b/src/Network/KRPC/Method.hs
@@ -8,11 +8,13 @@
 --   Normally, you don't need to import this module.
 --
 {-# LANGUAGE CPP                        #-}
-{-# LANGUAGE RankNTypes                 #-}
-{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE DefaultSignatures          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE RankNTypes                 #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
-{-# LANGUAGE DefaultSignatures          #-}
+{-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE FunctionalDependencies #-}
 module Network.KRPC.Method
        ( Method (..)
        , KRPC (..)
@@ -93,7 +95,9 @@ class ( Typeable req, Typeable resp
       , Serialize req, Serialize resp
 #endif
       )
-    => KRPC req resp where
+    => KRPC req resp | req -> resp, resp -> req where
+
+  type Envelope req resp
 
   -- | Method name. Default implementation uses lowercased @req@
   -- datatype name.
@@ -107,3 +111,6 @@ class ( Typeable req, Typeable resp
     where
       hole = error "krpc.method: impossible" :: req
 #endif
+
+  unseal :: Envelope req resp -> Either String req
+  seal :: resp -> Envelope req resp
diff --git a/src/Network/RPC.hs b/src/Network/RPC.hs
new file mode 100644
index 00000000..2e9356e8
--- /dev/null
+++ b/src/Network/RPC.hs
@@ -0,0 +1,15 @@
+{-# LANGUAGE RankNTypes #-}
+module Network.RPC where
+
+import Data.ByteString (ByteString)
+
+data MessageClass = Error | Query | Response
+ deriving (Eq,Ord,Enum,Bounded,Show,Read)
+
+data Messaging msg tid payload = Messaging
+    { messageClass       :: forall a. msg a -> MessageClass
+    , messageTransaction :: forall a. msg a -> tid
+    , messagePayload     :: forall a. msg a -> a
+    , encodePayload      :: msg payload -> msg ByteString
+    , decodePayload      :: msg ByteString -> Either String (msg payload)
+    }