4 files changed, 91 insertions, 26 deletions

diff --git a/src/Data/Torrent/InfoHash.hs b/src/Data/Torrent/InfoHash.hs
index 4d49fcb7..9eec631c 100644
--- a/src/Data/Torrent/InfoHash.hs
+++ b/src/Data/Torrent/InfoHash.hs
@@ -120,6 +120,7 @@ instance Convertible BS.ByteString InfoHash where
 -- | Parse infohash from base16\/base32\/base64 encoded string.
 instance Convertible Text InfoHash where
   safeConvert t
+      | 20 == hashLen = pure (InfoHash hashStr)
       | 26 <= hashLen && hashLen <= 28 =
         case Base64.decode hashStr of
           Left  msg   -> convError ("invalid base64 encoding " ++ msg) t
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 0733a85d..15d1099c 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -31,6 +31,7 @@ import Control.Applicative
 import Data.BEncode as BE
 import Data.BEncode.BDict
 import Data.ByteString as BS
+import Data.List as L
 import Data.Monoid
 import Data.Serialize as S
 import Data.Typeable
@@ -138,7 +139,7 @@ binary k = field (req k) >>= either (fail . format) return .
 
 instance (Typeable ip, Serialize ip) => BEncode (NodeFound ip) where
   toBEncode (NodeFound ns) = toDict $
-       nodes_key .=! S.encode ns
+       nodes_key .=! runPut (mapM_ put ns)
     .: endDict
 
   fromBEncode = fromDict $ NodeFound <$> binary nodes_key
@@ -154,7 +155,7 @@ instance (Serialize ip, Typeable ip)
 
 -- | Get peers associated with a torrent infohash.
 newtype GetPeers = GetPeers InfoHash
-  deriving Typeable
+  deriving (Show, Eq, Typeable)
 
 info_hash_key :: BKey
 info_hash_key = "info_hash"
@@ -172,7 +173,7 @@ data GotPeers ip = GotPeers
     -- | The token value is a required argument for a future
     -- announce_peer query.
   , grantedToken :: Token
-  } deriving Typeable
+  } deriving (Show, Eq, Typeable)
 
 peers_key :: BKey
 peers_key = "values"
@@ -182,17 +183,25 @@ 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
+    case peers of
+      Left  ns ->
+           nodes_key .=! runPut (mapM_ put ns)
+        .: token_key .=! grantedToken
+        .: endDict
+      Right ps ->
+           token_key .=! grantedToken
+        .: peers_key .=! L.map S.encode ps
+        .: endDict
 
-  fromBEncode = fromDict $ GotPeers <$> getPeerList  <*>! token_key
-    where
-      getPeerList = Right <$> binary peers_key
-                <|> Left  <$> binary nodes_key
+  fromBEncode = fromDict $ do
+     mns <- optional (binary nodes_key)                      -- "nodes"
+     tok <- field    (req    token_key)                      -- "token"
+     mps <- optional (field (req peers_key) >>= decodePeers) -- "values"
+     case (Right <$> mps) <|> (Left <$> mns) of
+       Nothing -> fail "get_peers: neihter peers nor nodes key is valid"
+       Just xs -> pure $ GotPeers xs tok
+     where
+       decodePeers = either fail pure . mapM S.decode
 
 instance (Typeable ip, Serialize ip) =>
          KRPC (Query GetPeers) (Response (GotPeers ip)) where
@@ -213,7 +222,7 @@ data Announce = Announce
 
     -- | received in response to a previous get_peers query.
   , sessionToken :: Token
-  } deriving Typeable
+  } deriving (Show, Eq, Typeable)
 
 port_key :: BKey
 port_key = "port"
@@ -235,6 +244,7 @@ instance BEncode Announce where
 -- port number under the infohash in its store of peer contact
 -- information.
 data Announced = Announced
+  deriving (Show, Eq, Typeable)
 
 instance BEncode Announced where
   toBEncode   _ = toBEncode Ping
diff --git a/tests/Network/BitTorrent/CoreSpec.hs b/tests/Network/BitTorrent/CoreSpec.hs
index 917ed5f5..9f9eb0eb 100644
--- a/tests/Network/BitTorrent/CoreSpec.hs
+++ b/tests/Network/BitTorrent/CoreSpec.hs
@@ -1,6 +1,9 @@
 -- | Re-export modules.
 module Network.BitTorrent.CoreSpec (spec) where
+import Network.BitTorrent.Core.FingerprintSpec as CoreSpec ()
+import Network.BitTorrent.Core.NodeSpec as CoreSpec ()
 import Network.BitTorrent.Core.PeerIdSpec as CoreSpec ()
+import Network.BitTorrent.Core.PeerAddrSpec as CoreSpec ()
 
 import Test.Hspec (Spec)
 
diff --git a/tests/Network/BitTorrent/DHT/MessageSpec.hs b/tests/Network/BitTorrent/DHT/MessageSpec.hs
index 9f4c58b0..52bda932 100644
--- a/tests/Network/BitTorrent/DHT/MessageSpec.hs
+++ b/tests/Network/BitTorrent/DHT/MessageSpec.hs
@@ -1,14 +1,19 @@
 {-# LANGUAGE RecordWildCards #-}
 module Network.BitTorrent.DHT.MessageSpec (spec) where
 import Control.Monad.Reader
-import Test.Hspec
 import Data.BEncode as BE
+import Data.ByteString.Lazy as BL
 import Data.Default
 import Data.List as L
 import Network.BitTorrent.Core
 import Network.BitTorrent.DHT.Message
 import Network.KRPC
 import Network.Socket (PortNumber)
+import Test.Hspec
+import Test.QuickCheck
+
+import Network.BitTorrent.CoreSpec ()
+import Data.Torrent.InfoHashSpec   ()
 
 
 remoteAddr :: SockAddr
@@ -29,6 +34,9 @@ rpc action = do
 isProtocolError :: KError -> Bool
 isProtocolError KError {..} = errorCode == ProtocolError
 
+prop_bencode :: Eq a => Show a => BEncode a => a -> Expectation
+prop_bencode x = BE.decode (BL.toStrict (BE.encode x)) `shouldBe` Right x
+
 spec :: Spec
 spec = do
  context ("you need running DHT node at " ++ show remoteAddr) $ do
@@ -46,8 +54,9 @@ spec = do
       BE.encode (Response "mnopqrstuvwxyz123456" Ping)
         `shouldBe` "d2:id20:mnopqrstuvwxyz123456e"
 
-    it "properly bencoded (iso)" $ do
-      pending
+    it "properly bencoded (iso)" $ property $ \ nid -> do
+      prop_bencode (Query    nid Ping)
+      prop_bencode (Response nid Ping)
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
@@ -74,8 +83,9 @@ spec = do
                 \e"
        `shouldBe` Right (Response nid (NodeFound [NodeInfo nid' naddr]))
 
-    it "properly bencoded (iso)" $ do
-      pending
+    it "properly bencoded (iso)" $ property $ \ nid x xs -> do
+      prop_bencode (Query    nid (FindNode  x))
+      prop_bencode (Response nid (NodeFound (xs :: [NodeInfo IPv4] )))
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
@@ -85,10 +95,38 @@ spec = do
 
   describe "get_peers" $ do
     it "properly bencoded" $ do
-      pending
-
-    it "properly bencoded (iso)" $ do
-      pending
+       BE.decode "d2:id20:abcdefghij0123456789\
+                  \9:info_hash20:mnopqrstuvwxyz123456\
+                  \e"
+         `shouldBe` Right (Query "abcdefghij0123456789"
+                           (GetPeers "mnopqrstuvwxyz123456"))
+
+       BE.decode "d2:id20:abcdefghij0123456789\
+                  \5:token8:aoeusnth\
+                  \6:valuesl6:\127\0\0\1\1\2\&6:\192\168\1\100\1\2e\
+                 \e"
+         `shouldBe` Right (Response "abcdefghij0123456789"
+             (GotPeers (Right [ "127.0.0.1:258" :: PeerAddr IPv4
+                              , "192.168.1.100:258"
+                              ]) "aoeusnth"))
+
+       BE.decode "d2:id20:abcdefghij0123456789\
+                  \5:nodes26:mnopqrstuvwxyz123456\127\0\0\1\1\2\
+                  \5:token8:aoeusnth\
+                 \e"
+         `shouldBe` Right (Response "abcdefghij0123456789"
+            (GotPeers
+             { peers = Left [NodeInfo "mnopqrstuvwxyz123456" "127.0.0.1:258"
+                                 :: NodeInfo IPv4]
+             , grantedToken = "aoeusnth"
+             }))
+
+    it "properly bencoded (iso)" $ property $ \ nid topic exs token -> do
+      prop_bencode (Query    nid (GetPeers topic))
+      let _ = exs :: Either [NodeInfo IPv4] [PeerAddr IPv4]
+      let nullPeerId paddr = paddr {peerId = Nothing}
+      let nullPeerIds      = either Left (Right . L.map nullPeerId)
+      prop_bencode (Response nid (GotPeers (nullPeerIds exs) token))
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
@@ -99,10 +137,23 @@ spec = do
 
   describe "announce" $ do
     it "properly bencoded" $ do
-      pending
+      BE.decode "d2:id20:abcdefghij0123456789\
+                 \9:info_hash20:mnopqrstuvwxyz123456\
+                 \4:porti6881e\
+                 \5:token8:aoeusnth\
+                \e" `shouldBe` Right
+        (Query "abcdefghij0123456789"
+                   (Announce "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
+
+
+      BE.decode "d2:id20:mnopqrstuvwxyz123456e"
+        `shouldBe` Right
+        (Response "mnopqrstuvwxyz123456" Announced)
+
+    it "properly bencoded (iso)" $ property $ \ nid topic port token -> do
+      prop_bencode (Query nid    (Announce topic port token))
+      prop_bencode (Response nid (Announced))
 
-    it "properly bencoded (iso)" $ do
-      pending
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId