4 files changed, 79 insertions, 9 deletions

diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index b2a95373..53cadcd8 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -26,6 +26,7 @@ library
                      , Data.Torrent.InfoHash
 
                      , Network.BitTorrent
+                     , Network.BitTorrent.Extension
 
                      , Network.BitTorrent.Peer
                      , Network.BitTorrent.Peer.Addr
diff --git a/src/Network/BitTorrent/Extension.hs b/src/Network/BitTorrent/Extension.hs
new file mode 100644
index 00000000..e37f3afb
--- /dev/null
+++ b/src/Network/BitTorrent/Extension.hs
@@ -0,0 +1,56 @@
+-- |
+--   Copyright   :  (c) Sam T. 2013
+--   License     :  MIT
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   This module provides peer capabilities detection.
+--
+--   > See http://www.bittorrent.org/beps/bep_0004.html
+--
+module Network.BitTorrent.Extension
+       ( Capabilities, ppCaps, defaultCaps, enabledCaps
+       , Extension, ppExtension, encodeExts, decodeExts
+       ) where
+
+import Data.Bits
+import Data.List
+import Data.Word
+
+
+type Capabilities = Word64
+
+ppCaps :: Capabilities -> String
+ppCaps = intercalate ", " . map ppExtension . decodeExts
+
+defaultCaps :: Capabilities
+defaultCaps = 0
+
+enabledCaps :: Capabilities -- ^ of the client.
+            -> Capabilities -- ^ of the peer.
+            -> Capabilities -- ^ should be considered as enabled.
+enabledCaps = (.&.)
+
+
+
+data Extension = ExtDHT  -- ^ BEP 5
+               | ExtFast -- ^ BEP 6
+                 deriving (Show, Eq, Ord, Enum, Bounded)
+
+ppExtension :: Extension -> String
+ppExtension ExtDHT  = "DHT"
+ppExtension ExtFast = "Fast Extension"
+
+extensionMask :: Extension -> Word64
+extensionMask ExtDHT  = 0x01
+extensionMask ExtFast = 0x04
+
+
+encodeExts :: [Extension] -> Capabilities
+encodeExts = foldr (.&.) 0 . map extensionMask
+
+decodeExts :: Capabilities -> [Extension]
+decodeExts rb = filter (testMask rb . extensionMask) [minBound..maxBound]
+  where
+    testMask bits x = bits .&. x > 0
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."