4 files changed, 252 insertions, 6 deletions

diff --git a/src/Network/BitTorrent/Core/Node.hs b/src/Network/BitTorrent/Core/Node.hs
index a1a87135..0cb95dd2 100644
--- a/src/Network/BitTorrent/Core/Node.hs
+++ b/src/Network/BitTorrent/Core/Node.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE TemplateHaskell            #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
 module Network.BitTorrent.Core.Node
        (  -- * Node ID
          NodeId
@@ -20,7 +21,9 @@ import Data.Aeson.TH
 import Data.Bits
 import Data.ByteString as BS
 import Data.BEncode as BE
+import Data.Default
 import Data.Ord
+import Data.Typeable
 import Data.Serialize as S
 import Data.Word
 import Network
@@ -37,11 +40,15 @@ import Network.BitTorrent.Core.PeerAddr ()
 -- | Normally, /this/ node id should we saved between invocations of
 -- the client software.
 newtype NodeId = NodeId ByteString
-  deriving (Show, Eq, Ord, BEncode, FromJSON, ToJSON)
+  deriving (Show, Eq, Ord, BEncode, FromJSON, ToJSON, Typeable)
 
 nodeIdSize :: Int
 nodeIdSize = 20
 
+-- | Meaningless node id, for testing purposes only.
+instance Default NodeId where
+  def = NodeId (BS.replicate nodeIdSize 0)
+
 instance Serialize NodeId where
   get = NodeId <$> getByteString nodeIdSize
   {-# INLINE get #-}
diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index bc4a1078..86b88491 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -254,7 +254,7 @@ peerSockAddr PeerAddr {..} =
 
 -- | Storage used to keep track a set of known peers in client,
 -- tracker or DHT sessions.
-newtype PeerStore a = PeerStore (HashMap InfoHash [PeerAddr a])
+newtype PeerStore ip = PeerStore (HashMap InfoHash [PeerAddr ip])
 
 -- | Empty store.
 instance Default (PeerStore a) where
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
new file mode 100644
index 00000000..85abf019
--- /dev/null
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -0,0 +1,240 @@
+-- | For more info see:
+--   <http://www.bittorrent.org/beps/bep_0005.html#dht-queries>
+--
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE FlexibleInstances  #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE UndecidableInstances  #-}
+module Network.BitTorrent.DHT.Message
+       ( -- * Envelopes
+         Query (..)
+       , Response (..)
+
+         -- * Queries
+         -- ** ping
+       , Ping (..)
+
+         -- ** find_node
+       , FindNode (..)
+       , NodeFound (..)
+
+         -- ** get_peers
+       , GetPeers (..)
+       , GotPeers (..)
+
+         -- ** announce_peer
+       , Announce  (..)
+       , Announced (..)
+       ) where
+
+import Control.Applicative
+import Data.BEncode as BE
+import Data.BEncode.BDict
+import Data.ByteString as BS
+import Data.Monoid
+import Data.Serialize as S
+import Data.Typeable
+import Network
+import Network.KRPC
+
+import Data.Torrent.InfoHash
+import Network.BitTorrent.Core
+import Network.KRPC ()
+
+{-----------------------------------------------------------------------
+-- envelopes
+-----------------------------------------------------------------------}
+
+type Token = ByteString
+
+node_id_key :: BKey
+node_id_key = "id"
+
+-- | All queries have an "id" key and value containing the node ID of
+-- the querying node.
+data Query a = Query
+  { thisNodeId  :: NodeId
+  , queryParams :: a
+  }
+
+instance BEncode a => BEncode (Query a) where
+  toBEncode Query {..} = toDict $
+      node_id_key .=! thisNodeId .: endDict
+      <>
+      dict (toBEncode queryParams)
+    where
+      dict (BDict d) = d
+      dict    _      = error "impossible: instance BEncode (Query a)"
+
+  fromBEncode v = do
+    Query <$> fromDict (field (req node_id_key)) v
+          <*> fromBEncode v
+
+-- | All responses have an "id" key and value containing the node ID
+-- of the responding node.
+data Response a = Response
+  { remoteNodeId :: NodeId
+  , responseVals :: a
+  }
+
+instance BEncode a => BEncode (Response a) where
+  toBEncode = toBEncode . toQuery
+    where
+      toQuery (Response nid a) = Query nid a
+
+  fromBEncode b = fromQuery <$> fromBEncode b
+    where
+      fromQuery (Query nid a) = Response nid a
+
+
+{-----------------------------------------------------------------------
+-- ping method
+-----------------------------------------------------------------------}
+
+-- | The most basic query is a ping.
+data Ping = Ping
+  deriving Typeable
+
+instance BEncode Ping where
+  toBEncode Ping = toDict endDict
+  fromBEncode _  = pure Ping
+
+-- | \"q\" = \"ping\"
+instance KRPC (Query Ping) [Ping] where
+  method = "ping"
+
+{-----------------------------------------------------------------------
+-- find_node method
+-----------------------------------------------------------------------}
+
+-- | Find node is used to find the contact information for a node
+-- given its ID.
+newtype FindNode = FindNode NodeId
+  deriving Typeable
+
+target_key :: BKey
+target_key = "target"
+
+instance BEncode FindNode where
+  toBEncode (FindNode nid) = toDict   $ target_key .=! nid .: endDict
+  fromBEncode              = fromDict $ FindNode  <$>! target_key
+
+-- | When a node receives a find_node query, it should respond with a
+-- the compact node info for the target node or the K (8) closest good
+-- nodes in its own routing table.
+--
+newtype NodeFound ip = NodeFound [NodeInfo ip]
+  deriving Typeable
+
+nodes_key :: BKey
+nodes_key = "nodes"
+
+binary :: Serialize a => BE.Get BS.ByteString -> BE.Get a
+binary m = m >>= either fail return . S.decode
+
+instance (Typeable ip, Serialize ip) => BEncode (NodeFound ip) where
+  toBEncode (NodeFound ns) = toDict $
+       nodes_key .=! S.encode ns
+    .: endDict
+
+  fromBEncode = fromDict $ NodeFound <$> binary (field (req nodes_key))
+
+-- | \"q\" == \"find_node\"
+instance (Serialize ip, Typeable ip)
+      => KRPC (Query FindNode) (Response (NodeFound ip)) where
+  method = "find_node"
+
+{-----------------------------------------------------------------------
+-- get_peers method
+-----------------------------------------------------------------------}
+
+-- | Get peers associated with a torrent infohash.
+newtype GetPeers = GetPeers InfoHash
+  deriving Typeable
+
+info_hash_key :: BKey
+info_hash_key = "info_hash"
+
+instance BEncode GetPeers where
+  toBEncode (GetPeers ih) = toDict   $ info_hash_key .=! ih .: endDict
+  fromBEncode             = fromDict $ GetPeers <$>! info_hash_key
+
+data GotPeers ip = GotPeers
+  { -- | If the queried node has no peers for the infohash, returned
+    -- the K nodes in the queried nodes routing table closest to the
+    -- infohash supplied in the query.
+    peers        :: Either [NodeInfo ip] [PeerAddr ip]
+
+    -- | The token value is a required argument for a future
+    -- announce_peer query.
+  , grantedToken :: Token
+  } deriving Typeable
+
+peers_key :: BKey
+peers_key = "values"
+
+token_key :: BKey
+token_key = "token"
+
+instance (Typeable ip, Serialize ip) => BEncode (GotPeers ip) where
+  toBEncode GotPeers {..} = toDict $
+       putPeerList peers
+    .: token_key .=! grantedToken
+    .: endDict
+    where
+      putPeerList (Right ps) = peers_key .=! S.encode ps
+      putPeerList (Left  ns) = nodes_key .=! S.encode ns
+
+  fromBEncode = fromDict $ GotPeers <$> getPeerList  <*>! token_key
+    where
+      getPeerList = Right <$> binary (field (req peers_key))
+                <|> Left  <$> binary (field (req nodes_key))
+
+instance (Typeable ip, Serialize ip) =>
+         KRPC (Query GetPeers) (Response (GotPeers ip)) where
+  method = "get_peers"
+
+{-----------------------------------------------------------------------
+-- announce method
+-----------------------------------------------------------------------}
+
+-- | Announce that the peer, controlling the querying node, is
+-- downloading a torrent on a port.
+data Announce = Announce
+  { -- | infohash of the torrent;
+    topic    :: InfoHash
+
+    -- | the port /this/ peer is listening;
+  , port     :: PortNumber
+
+    -- | received in response to a previous get_peers query.
+  , sessionToken :: Token
+  } deriving Typeable
+
+port_key :: BKey
+port_key = "port"
+
+instance BEncode Announce where
+  toBEncode Announce {..} = toDict $
+       info_hash_key .=! topic
+    .: port_key      .=! port
+    .: token_key     .=! sessionToken
+    .: endDict
+  fromBEncode = fromDict $ do
+    Announce <$>! info_hash_key
+             <*>! port_key
+             <*>! token_key
+
+-- | The queried node must verify that the token was previously sent
+-- to the same IP address as the querying node. Then the queried node
+-- should store the IP address of the querying node and the supplied
+-- port number under the infohash in its store of peer contact
+-- information.
+data Announced = Announced
+
+instance BEncode Announced where
+  toBEncode   _ = toBEncode Ping
+  fromBEncode _ = pure  Announced
+
+instance KRPC (Query Announce) (Response Announced) where
+  method = "announce_peer"
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index 8d0c62f4..e93f8bbe 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -679,8 +679,7 @@ instance Default ExtendedHandshake where
 
 instance Monoid ExtendedHandshake where
     mempty = def { ehsCaps = mempty }
-    mappend old new =
-        ExtendedHandshake {
+    mappend old new = ExtendedHandshake {
       ehsCaps         = ehsCaps old <> ehsCaps new,
       ehsIPv4         = ehsIPv4 old         `mergeOld` ehsIPv4 new,
       ehsIPv6         = ehsIPv6 old         `mergeOld` ehsIPv6 new,
@@ -691,8 +690,8 @@ instance Monoid ExtendedHandshake where
       ehsYourIp       = ehsYourIp old       `mergeOld` ehsYourIp new
     }
       where
-        mergeOld old new = old <|> new
-        mergeNew old new = new <|> old
+        mergeOld mold mnew = mold <|> mnew
+        mergeNew mold mnew = mnew <|> mold
 
 
 instance BEncode ExtendedHandshake where