2 files changed, 22 insertions, 9 deletions

diff --git a/src/Network/BitTorrent/PeerWire/Block.hs b/src/Network/BitTorrent/PeerWire/Block.hs
index 582accdb..fbc65338 100644
--- a/src/Network/BitTorrent/PeerWire/Block.hs
+++ b/src/Network/BitTorrent/PeerWire/Block.hs
@@ -1,5 +1,7 @@
 module Network.BitTorrent.PeerWire.Block
-       ( BlockIx(..), Block(..), PieceIx
+       ( BlockIx(..)
+       , Block(..), blockSize
+       , PieceIx
        , BlockLIx, PieceLIx
        , defaultBlockSize
        , pieceIx, blockIx
@@ -68,6 +70,9 @@ data Block = Block {
 ppBlock :: Block -> String
 ppBlock = ppBlockIx . blockIx
 
+blockSize :: Block -> Int
+blockSize blk = B.length (blkData blk)
+
 -- | Widely used semi-official block size.
 defaultBlockSize :: Int
 defaultBlockSize = 16 * 1024
diff --git a/src/Network/BitTorrent/PeerWire/Handshake.hs b/src/Network/BitTorrent/PeerWire/Handshake.hs
index e0d1672b..62d7d7f4 100644
--- a/src/Network/BitTorrent/PeerWire/Handshake.hs
+++ b/src/Network/BitTorrent/PeerWire/Handshake.hs
@@ -12,7 +12,7 @@
 --
 {-# LANGUAGE OverloadedStrings #-}
 module Network.BitTorrent.PeerWire.Handshake
-       ( Handshake
+       ( Handshake, handshakeCaps
        , handshake
        , ppHandshake
        , defaultHandshake, defaultBTProtocol, defaultReserved
@@ -29,6 +29,7 @@ import Data.Torrent.InfoHash
 import Network
 import Network.Socket.ByteString
 
+import Network.BitTorrent.Extension
 import Network.BitTorrent.Peer.ID
 import Network.BitTorrent.Peer.ClientInfo
 
@@ -69,6 +70,10 @@ instance Serialize Handshake where
               <*> get
               <*> get
 
+
+handshakeCaps :: Handshake -> Capabilities
+handshakeCaps = hsReserved
+
 -- TODO add reserved bits info
 -- | Format handshake in human readable form.
 ppHandshake :: Handshake -> String
@@ -95,7 +100,7 @@ defaultReserved = 0
 defaultHandshake :: InfoHash -> PeerID -> Handshake
 defaultHandshake = Handshake defaultBTProtocol defaultReserved
 
-
+-- TODO exceptions instead of Either
 -- | Handshaking with a peer specified by the second argument.
 --
 handshake :: Socket -> Handshake -> IO (Either String Handshake)
@@ -103,12 +108,15 @@ handshake sock hs = do
     sendAll sock (S.encode hs)
 
     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))
+    if B.length header == 0 then
+        return $ Left ""
+      else do
+        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."