2 files changed, 17 insertions, 4 deletions

diff --git a/src/Network/BitTorrent/PeerWire/Handshake.hs b/src/Network/BitTorrent/PeerWire/Handshake.hs
index e4d27f75..e0d1672b 100644
--- a/src/Network/BitTorrent/PeerWire/Handshake.hs
+++ b/src/Network/BitTorrent/PeerWire/Handshake.hs
@@ -75,9 +75,13 @@ ppHandshake :: Handshake -> String
 ppHandshake hs = BC.unpack (hsProtocol hs) ++ " "
               ++ ppClientInfo (clientInfo (hsPeerID hs))
 
+-- | Get handshake message size in bytes from the length of protocol string.
+handshakeSize :: Word8 -> Int
+handshakeSize n = 1 + fromIntegral n + 8 + 20 + 20
+
 -- | Maximum size of handshake message in bytes.
 handshakeMaxSize :: Int
-handshakeMaxSize = 1 + 256 + 8 + 20 + 20
+handshakeMaxSize = handshakeSize 255
 
 -- | Default protocol string "BitTorrent protocol" as is.
 defaultBTProtocol :: ByteString
@@ -97,8 +101,14 @@ defaultHandshake = Handshake defaultBTProtocol defaultReserved
 handshake :: Socket -> Handshake -> IO (Either String Handshake)
 handshake sock hs = do
     sendAll sock (S.encode hs)
-    r <- recv sock handshakeMaxSize
-    return (checkIH (S.decode r))
+
+    header <- recv sock 1
+    let protocolLen = B.head header
+    let restLen     = handshakeSize protocolLen - 1
+    body <- recv sock restLen
+    let resp = B.cons protocolLen body
+
+    return (checkIH (S.decode resp))
   where
     checkIH (Right hs')
       | hsInfoHash hs /= hsInfoHash hs' = Left "Handshake info hash do not match."
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 39102eed..1c520221 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -99,7 +99,10 @@ instance Serialize Message where
           0x0D -> SuggestPiece  <$> getInt
           0x10 -> RejectRequest <$> get
           0x11 -> AllowedFast   <$> getInt
-          _    -> fail $ "unknown message ID: " ++ show mid
+          _    -> do
+            rm <- remaining >>= getBytes
+            fail $ "unknown message ID: " ++ show mid ++ "\n"
+                ++ "remaining available bytes: " ++ show rm
 
     where
       getBlock :: Int -> Get Block