1 files changed, 1232 insertions, 0 deletions

diff --git a/bittorrent/src/Network/BitTorrent/Exchange/Message.hs b/bittorrent/src/Network/BitTorrent/Exchange/Message.hs
new file mode 100644
index 00000000..2c6770f7
--- /dev/null
+++ b/bittorrent/src/Network/BitTorrent/Exchange/Message.hs
@@ -0,0 +1,1232 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   Normally peer to peer communication consisting of the following
+--   steps:
+--
+--   * In order to establish the connection between peers we should
+--   send 'Handshake' message. The 'Handshake' is a required message
+--   and must be the first message transmitted by the peer to the
+--   another peer. Another peer should reply with a handshake as well.
+--
+--   * Next peer might sent bitfield message, but might not. In the
+--   former case we should update bitfield peer have. Again, if we
+--   have some pieces we should send bitfield. Normally bitfield
+--   message should sent after the handshake message.
+--
+--   * Regular exchange messages. TODO docs
+--
+--   For more high level API see "Network.BitTorrent.Exchange" module.
+--
+--   For more infomation see:
+--   <https://wiki.theory.org/BitTorrentSpecification#Peer_wire_protocol_.28TCP.29>
+--
+{-# LANGUAGE ViewPatterns               #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE FlexibleContexts           #-}
+{-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE TemplateHaskell            #-}
+{-# OPTIONS  -fno-warn-orphans          #-}
+module Network.BitTorrent.Exchange.Message
+       ( -- * Capabilities
+         Capabilities (..)
+       , Extension (..)
+       , Caps
+
+         -- * Handshake
+       , ProtocolName
+       , Handshake(..)
+       , defaultHandshake
+       , handshakeSize
+       , handshakeMaxSize
+       , handshakeStats
+
+         -- * Stats
+       , ByteCount
+       , ByteStats   (..)
+       , byteLength
+
+         -- * Messages
+       , Message        (..)
+       , defaultKeepAliveTimeout
+       , defaultKeepAliveInterval
+       , PeerMessage    (..)
+
+         -- ** Core messages
+       , StatusUpdate   (..)
+       , Available      (..)
+       , Transfer       (..)
+       , defaultRequestQueueLength
+
+         -- ** Fast extension
+       , FastMessage    (..)
+
+         -- ** Extension protocol
+       , ExtendedMessage   (..)
+
+         -- *** Capabilities
+       , ExtendedExtension (..)
+       , ExtendedCaps      (..)
+
+         -- *** Handshake
+       , ExtendedHandshake (..)
+       , defaultQueueLength
+       , nullExtendedHandshake
+
+         -- *** Metadata
+       , ExtendedMetadata  (..)
+       , metadataPieceSize
+       , defaultMetadataFactor
+       , defaultMaxInfoDictSize
+       , isLastPiece
+       , isValidPiece
+       ) where
+
+import Control.Applicative
+import Control.Arrow ((&&&), (***))
+import Control.Monad (when)
+import Data.Attoparsec.ByteString.Char8 as BS
+import Data.BEncode as BE
+import Data.BEncode.BDict as BE
+import Data.BEncode.Internal as BE (ppBEncode, parser)
+import Data.BEncode.Types    (BDict)
+import Data.Bits
+import Data.ByteString as BS
+import Data.ByteString.Char8 as BC
+import Data.ByteString.Lazy  as BL
+import Data.Default
+import Data.List as L
+import Data.Map.Strict as M
+import Data.Maybe
+import Data.Monoid
+import Data.Ord
+import Data.Serialize as S
+import Data.String
+import Data.Text as T
+import Data.Typeable
+import Data.Word
+import Data.IP
+import Network
+import Network.Socket hiding (KeepAlive)
+import Text.PrettyPrint as PP hiding ((<>))
+import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
+
+import Data.Torrent hiding (Piece (..))
+import qualified Data.Torrent as P (Piece (..))
+import Network.Address
+import Network.BitTorrent.Exchange.Bitfield
+import Network.BitTorrent.Exchange.Block
+
+{-----------------------------------------------------------------------
+--  Capabilities
+-----------------------------------------------------------------------}
+
+-- |
+class Capabilities caps where
+  type Ext caps :: *
+
+  -- | Pack extensions to caps.
+  toCaps   :: [Ext caps] -> caps
+
+  -- | Unpack extensions from caps.
+  fromCaps :: caps -> [Ext caps]
+
+  -- | Check if an extension is a member of the specified set.
+  allowed  :: Ext caps -> caps -> Bool
+
+ppCaps :: Capabilities caps => Pretty (Ext caps) => caps -> Doc
+ppCaps = hcat . punctuate ", " . L.map pPrint . fromCaps
+
+{-----------------------------------------------------------------------
+--  Extensions
+-----------------------------------------------------------------------}
+
+-- | Enumeration of message extension protocols.
+--
+--   For more info see: <http://www.bittorrent.org/beps/bep_0004.html>
+--
+data Extension
+  = ExtDHT      -- ^ BEP 5:  allow to send PORT messages.
+  | ExtFast     -- ^ BEP 6:  allow to send FAST messages.
+  | ExtExtended -- ^ BEP 10: allow to send the extension protocol messages.
+    deriving (Show, Eq, Ord, Enum, Bounded)
+
+-- | Full extension names, suitable for logging.
+instance Pretty Extension where
+  pPrint ExtDHT      = "Distributed Hash Table Protocol"
+  pPrint ExtFast     = "Fast Extension"
+  pPrint ExtExtended = "Extension Protocol"
+
+-- | Extension bitmask as specified by BEP 4.
+extMask :: Extension -> Word64
+extMask ExtDHT      = 0x01
+extMask ExtFast     = 0x04
+extMask ExtExtended = 0x100000
+
+{-----------------------------------------------------------------------
+--  Capabilities
+-----------------------------------------------------------------------}
+
+-- | Capabilities is a set of 'Extension's usually sent in 'Handshake'
+-- messages.
+newtype Caps = Caps Word64
+  deriving (Show, Eq)
+
+-- | Render set of extensions as comma separated list.
+instance Pretty Caps where
+  pPrint = ppCaps
+  {-# INLINE pPrint #-}
+
+-- | The empty set.
+instance Default Caps where
+  def = Caps 0
+  {-# INLINE def #-}
+
+-- | Monoid under intersection. 'mempty' includes all known extensions.
+instance Monoid Caps where
+  mempty = toCaps [minBound .. maxBound]
+  {-# INLINE mempty #-}
+
+  mappend (Caps a) (Caps b) = Caps (a .&. b)
+  {-# INLINE mappend #-}
+
+-- | 'Handshake' compatible encoding.
+instance Serialize Caps where
+  put (Caps caps) = S.putWord64be caps
+  {-# INLINE put #-}
+
+  get = Caps <$> S.getWord64be
+  {-# INLINE get #-}
+
+instance Capabilities Caps where
+  type Ext Caps = Extension
+
+  allowed e (Caps caps) = (extMask e .&. caps) /= 0
+  {-# INLINE allowed #-}
+
+  toCaps        = Caps . L.foldr (.|.) 0 . L.map extMask
+  fromCaps caps = L.filter (`allowed` caps) [minBound..maxBound]
+
+{-----------------------------------------------------------------------
+    Handshake
+-----------------------------------------------------------------------}
+
+maxProtocolNameSize :: Word8
+maxProtocolNameSize = maxBound
+
+-- | The protocol name is used to identify to the local peer which
+-- version of BTP the remote peer uses.
+newtype ProtocolName = ProtocolName BS.ByteString
+  deriving (Eq, Ord, Typeable)
+
+-- | In BTP/1.0 the name is 'BitTorrent protocol'. If this string is
+-- different from the local peers own protocol name, then the
+-- connection is to be dropped.
+instance Default ProtocolName where
+  def = ProtocolName "BitTorrent protocol"
+
+instance Show ProtocolName where
+  show (ProtocolName bs) = show bs
+
+instance Pretty ProtocolName where
+  pPrint (ProtocolName bs) = PP.text $ BC.unpack bs
+
+instance IsString ProtocolName where
+  fromString str
+    | L.length str <= fromIntegral maxProtocolNameSize
+    = ProtocolName (fromString str)
+    | otherwise = error $ "fromString: ProtocolName too long: " ++ str
+
+instance Serialize ProtocolName where
+  put (ProtocolName bs) = do
+    putWord8 $ fromIntegral $ BS.length bs
+    putByteString bs
+
+  get = do
+    len  <- getWord8
+    bs   <- getByteString $ fromIntegral len
+    return (ProtocolName bs)
+
+-- | Handshake message is used to exchange all information necessary
+-- to establish connection between peers.
+--
+data Handshake = Handshake {
+    -- | Identifier of the protocol. This is usually equal to 'def'.
+    hsProtocol    :: ProtocolName
+
+    -- | Reserved bytes used to specify supported BEP's.
+  , hsReserved    :: Caps
+
+    -- | Info hash of the info part of the metainfo file. that is
+    -- transmitted in tracker requests. Info hash of the initiator
+    -- handshake and response handshake should match, otherwise
+    -- initiator should break the connection.
+    --
+  , hsInfoHash    :: InfoHash
+
+    -- | Peer id of the initiator. This is usually the same peer id
+    -- that is transmitted in tracker requests.
+    --
+  , hsPeerId      :: PeerId
+
+  } deriving (Show, Eq)
+
+instance Serialize Handshake where
+  put Handshake {..} = do
+    put hsProtocol
+    put hsReserved
+    put hsInfoHash
+    put hsPeerId
+  get = Handshake <$> get <*> get <*> get <*> get
+
+-- | Show handshake protocol string, caps and fingerprint.
+instance Pretty Handshake where
+  pPrint Handshake {..}
+    = pPrint hsProtocol           $$
+      pPrint hsReserved           $$
+      pPrint (fingerprint hsPeerId)
+
+-- | 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 = handshakeSize maxProtocolNameSize
+
+-- | Handshake with default protocol string and reserved bitmask.
+defaultHandshake :: InfoHash -> PeerId -> Handshake
+defaultHandshake = Handshake def def
+
+handshakeStats :: Handshake -> ByteStats
+handshakeStats (Handshake (ProtocolName bs) _ _ _)
+  = ByteStats 1 (BS.length bs + 8 + 20 + 20) 0
+
+{-----------------------------------------------------------------------
+--  Stats
+-----------------------------------------------------------------------}
+
+-- | Number of bytes.
+type ByteCount = Int
+
+-- | Summary of encoded message byte layout can be used to collect
+-- stats about message flow in both directions. This data can be
+-- retrieved using 'stats' function.
+data ByteStats = ByteStats
+  { -- | Number of bytes used to help encode 'control' and 'payload'
+    -- bytes: message size, message ID's, etc
+    overhead :: {-# UNPACK #-} !ByteCount
+
+    -- | Number of bytes used to exchange peers state\/options: piece
+    -- and block indexes, infohash, port numbers, peer ID\/IP, etc.
+  , control  :: {-# UNPACK #-} !ByteCount
+
+    -- | Number of payload bytes: torrent data blocks and infodict
+    -- metadata.
+  , payload  :: {-# UNPACK #-} !ByteCount
+  } deriving Show
+
+instance Pretty ByteStats where
+  pPrint s @ ByteStats {..} = fsep
+    [ PP.int overhead, "overhead"
+    , PP.int control,  "control"
+    , PP.int payload,  "payload"
+    , "bytes"
+    ] $+$ fsep
+    [ PP.int (byteLength s), "total bytes"
+    ]
+
+-- | Empty byte sequences.
+instance Default ByteStats where
+  def = ByteStats 0 0 0
+
+-- | Monoid under addition.
+instance Monoid ByteStats where
+  mempty      = def
+  mappend a b = ByteStats
+    { overhead = overhead a + overhead b
+    , control  = control  a + control  b
+    , payload  = payload  a + payload  b
+    }
+
+-- | Sum of the all byte sequences.
+byteLength :: ByteStats -> Int
+byteLength ByteStats {..} = overhead + control + payload
+
+{-----------------------------------------------------------------------
+--  Regular messages
+-----------------------------------------------------------------------}
+
+-- | Messages which can be sent after handshaking. Minimal complete
+-- definition: 'envelop'.
+class PeerMessage a where
+  -- | Construct a message to be /sent/. Note that if 'ExtendedCaps'
+  -- do not contain mapping for this message the default
+  -- 'ExtendedMessageId' is used.
+  envelop :: ExtendedCaps -- ^ The /receiver/ extended capabilities;
+          -> a            -- ^ An regular message;
+          -> Message      -- ^ Enveloped message to sent.
+
+  -- | Find out the extension this message belong to. Can be used to
+  -- check if this message is allowed to send\/recv in current
+  -- session.
+  requires :: a -> Maybe Extension
+  requires _ = Nothing
+
+  -- | Get sizes of overhead\/control\/payload byte sequences of
+  -- binary message representation without encoding message to binary
+  -- bytestring.
+  --
+  --   This function should obey one law:
+  --
+  --     * 'byteLength' ('stats' msg) == 'BL.length' ('encode' msg)
+  --
+  stats :: a -> ByteStats
+  stats _ = ByteStats 4 0 0
+
+{-----------------------------------------------------------------------
+--  Status messages
+-----------------------------------------------------------------------}
+
+-- | Notification that the sender have updated its
+-- 'Network.BitTorrent.Exchange.Status.PeerStatus'.
+data StatusUpdate
+    -- | Notification that the sender will not upload data to the
+    -- receiver until unchoking happen.
+  = Choking    !Bool
+
+    -- | Notification that the sender is interested (or not interested)
+    -- in any of the receiver's data pieces.
+  | Interested !Bool
+    deriving (Show, Eq, Ord, Typeable)
+
+instance Pretty StatusUpdate where
+  pPrint (Choking    False) = "not choking"
+  pPrint (Choking    True ) = "choking"
+  pPrint (Interested False) = "not interested"
+  pPrint (Interested True ) = "interested"
+
+instance PeerMessage StatusUpdate where
+  envelop _ = Status
+  {-# INLINE envelop #-}
+
+  stats _ = ByteStats 4 1 0
+  {-# INLINE stats #-}
+
+{-----------------------------------------------------------------------
+--  Available messages
+-----------------------------------------------------------------------}
+
+-- | Messages used to inform receiver which pieces of the torrent
+-- sender have.
+data Available =
+    -- | Zero-based index of a piece that has just been successfully
+    -- downloaded and verified via the hash.
+    Have    ! PieceIx
+
+    -- | The bitfield message may only be sent immediately after the
+    -- handshaking sequence is complete, and before any other message
+    -- are sent. If client have no pieces then bitfield need not to be
+    -- sent.
+  | Bitfield !Bitfield
+    deriving (Show, Eq)
+
+instance Pretty Available where
+  pPrint (Have     ix ) = "Have"     <+> int ix
+  pPrint (Bitfield _  ) = "Bitfield"
+
+instance PeerMessage Available where
+  envelop _ = Available
+  {-# INLINE envelop #-}
+
+  stats (Have      _) = ByteStats (4 + 1) 4 0
+  stats (Bitfield bf) = ByteStats (4 + 1) (q + trailing)  0
+    where
+      trailing = if r == 0 then 0 else 1
+      (q, r) = quotRem (totalCount bf) 8
+
+{-----------------------------------------------------------------------
+--  Transfer messages
+-----------------------------------------------------------------------}
+
+-- | Messages used to transfer 'Block's.
+data Transfer
+    -- | Request for a particular block. If a client is requested a
+    -- block that another peer do not have the peer might not answer
+    -- at all.
+  = Request ! BlockIx
+
+    -- | Response to a request for a block.
+  | Piece   !(Block BL.ByteString)
+
+    -- | Used to cancel block requests. It is typically used during
+    -- "End Game".
+  | Cancel  !BlockIx
+    deriving (Show, Eq)
+
+instance Pretty Transfer where
+  pPrint (Request  ix ) = "Request"  <+> pPrint ix
+  pPrint (Piece    blk) = "Piece"    <+> pPrint blk
+  pPrint (Cancel   i  ) = "Cancel"   <+> pPrint i
+
+instance PeerMessage Transfer where
+  envelop _ = Transfer
+  {-# INLINE envelop #-}
+
+  stats (Request _  ) = ByteStats (4 + 1) (3 * 4) 0
+  stats (Piece   p  ) = ByteStats (4 + 1) (4 + 4 + blockSize p) 0
+  stats (Cancel  _  ) = ByteStats (4 + 1) (3 * 4) 0
+
+-- TODO increase
+-- | Max number of pending 'Request's inflight.
+defaultRequestQueueLength :: Int
+defaultRequestQueueLength = 1
+
+{-----------------------------------------------------------------------
+--  Fast messages
+-----------------------------------------------------------------------}
+
+-- | BEP6 messages.
+data FastMessage =
+    -- | If a peer have all pieces it might send the 'HaveAll' message
+    -- instead of 'Bitfield' message. Used to save bandwidth.
+    HaveAll
+
+    -- | If a peer have no pieces it might send 'HaveNone' message
+    -- intead of 'Bitfield' message. Used to save bandwidth.
+  | HaveNone
+
+    -- | This is an advisory message meaning "you might like to
+    -- download this piece." Used to avoid excessive disk seeks and
+    -- amount of IO.
+  | SuggestPiece  !PieceIx
+
+    -- | Notifies a requesting peer that its request will not be
+    -- satisfied.
+  | RejectRequest !BlockIx
+
+    -- | This is an advisory messsage meaning \"if you ask for this
+    -- piece, I'll give it to you even if you're choked.\" Used to
+    -- shorten starting phase.
+  | AllowedFast   !PieceIx
+    deriving (Show, Eq)
+
+instance Pretty FastMessage where
+  pPrint (HaveAll          ) = "Have all"
+  pPrint (HaveNone         ) = "Have none"
+  pPrint (SuggestPiece  pix) = "Suggest"      <+> int    pix
+  pPrint (RejectRequest bix) = "Reject"       <+> pPrint bix
+  pPrint (AllowedFast   pix) = "Allowed fast" <+> int    pix
+
+instance PeerMessage FastMessage where
+  envelop  _ = Fast
+  {-# INLINE envelop #-}
+
+  requires _ = Just ExtFast
+  {-# INLINE requires #-}
+
+  stats  HaveAll          = ByteStats 4 1  0
+  stats  HaveNone         = ByteStats 4 1  0
+  stats (SuggestPiece  _) = ByteStats 5 4  0
+  stats (RejectRequest _) = ByteStats 5 12 0
+  stats (AllowedFast   _) = ByteStats 5 4  0
+
+{-----------------------------------------------------------------------
+--  Extension protocol
+-----------------------------------------------------------------------}
+
+{-----------------------------------------------------------------------
+--  Extended capabilities
+-----------------------------------------------------------------------}
+
+data ExtendedExtension
+  = ExtMetadata -- ^ BEP 9: Extension for Peers to Send Metadata Files
+    deriving (Show, Eq, Ord, Enum, Bounded, Typeable)
+
+instance IsString ExtendedExtension where
+  fromString = fromMaybe (error msg) . fromKey . fromString
+    where
+      msg = "fromString: could not parse ExtendedExtension"
+
+instance Pretty ExtendedExtension where
+  pPrint ExtMetadata = "Extension for Peers to Send Metadata Files"
+
+fromKey :: BKey -> Maybe ExtendedExtension
+fromKey "ut_metadata" = Just ExtMetadata
+fromKey _             = Nothing
+{-# INLINE fromKey #-}
+
+toKey :: ExtendedExtension -> BKey
+toKey ExtMetadata = "ut_metadata"
+{-# INLINE toKey #-}
+
+type ExtendedMessageId = Word8
+
+extId :: ExtendedExtension -> ExtendedMessageId
+extId ExtMetadata = 1
+{-# INLINE extId #-}
+
+type ExtendedMap = Map ExtendedExtension ExtendedMessageId
+
+-- | The extension IDs must be stored for every peer, because every
+-- peer may have different IDs for the same extension.
+--
+newtype ExtendedCaps = ExtendedCaps { extendedCaps :: ExtendedMap }
+  deriving (Show, Eq)
+
+instance Pretty ExtendedCaps where
+  pPrint = ppCaps
+  {-# INLINE pPrint #-}
+
+-- | The empty set.
+instance Default ExtendedCaps where
+  def = ExtendedCaps M.empty
+
+-- | Monoid under intersection:
+--
+--     * The 'mempty' caps includes all known extensions;
+--
+--     * the 'mappend' operation is NOT commutative: it return message
+-- id from the first caps for the extensions existing in both caps.
+--
+instance Monoid ExtendedCaps where
+  mempty  = toCaps [minBound..maxBound]
+  mappend (ExtendedCaps a) (ExtendedCaps b) =
+    ExtendedCaps (M.intersection a b)
+
+appendBDict ::  BDict -> ExtendedMap -> ExtendedMap
+appendBDict (Cons key val xs) caps
+  | Just  ext <- fromKey     key
+  , Right eid <- fromBEncode val   = M.insert ext eid (appendBDict xs caps)
+  |           otherwise            = appendBDict xs caps
+appendBDict  Nil              caps = caps
+
+-- | Handshake compatible encoding.
+instance BEncode ExtendedCaps where
+  toBEncode = BDict . BE.fromAscList . L.sortBy (comparing fst)
+            . L.map (toKey *** toBEncode) . M.toList . extendedCaps
+
+  fromBEncode (BDict bd) = pure $ ExtendedCaps $ appendBDict bd M.empty
+  fromBEncode _          = decodingError "ExtendedCaps"
+
+instance Capabilities ExtendedCaps where
+  type Ext ExtendedCaps = ExtendedExtension
+
+  toCaps = ExtendedCaps . M.fromList . L.map (id &&& extId)
+
+  fromCaps = M.keys . extendedCaps
+  {-# INLINE fromCaps #-}
+
+  allowed e (ExtendedCaps caps) = M.member e caps
+  {-# INLINE allowed #-}
+
+remoteMessageId :: ExtendedExtension -> ExtendedCaps -> ExtendedMessageId
+remoteMessageId ext = fromMaybe (extId ext) . M.lookup ext . extendedCaps
+
+{-----------------------------------------------------------------------
+--  Extended handshake
+-----------------------------------------------------------------------}
+
+-- | This message should be sent immediately after the standard
+-- bittorrent handshake to any peer that supports this extension
+-- protocol. Extended handshakes can be sent more than once, however
+-- an implementation may choose to ignore subsequent handshake
+-- messages.
+--
+data ExtendedHandshake = ExtendedHandshake
+  { -- | If this peer has an IPv4 interface, this is the compact
+    -- representation of that address.
+    ehsIPv4        :: Maybe HostAddress
+
+    -- | If this peer has an IPv6 interface, this is the compact
+    -- representation of that address.
+  , ehsIPv6        :: Maybe HostAddress6
+
+    -- | Dictionary of supported extension messages which maps names
+    -- of extensions to an extended message ID for each extension
+    -- message.
+  , ehsCaps        :: ExtendedCaps
+
+    -- | Size of 'Data.Torrent.InfoDict' in bytes. This field should
+    -- be added if 'ExtMetadata' is enabled in current session /and/
+    -- peer have the torrent file.
+  , ehsMetadataSize :: Maybe Int
+
+    -- | Local TCP /listen/ port. Allows each side to learn about the
+    -- TCP port number of the other side.
+  , ehsPort        :: Maybe PortNumber
+
+    -- | Request queue the number of outstanding 'Request' messages
+    -- this client supports without dropping any.
+  , ehsQueueLength :: Maybe Int
+
+    -- | Client name and version.
+  , ehsVersion     :: Maybe Text
+
+    -- | IP of the remote end
+  , ehsYourIp      :: Maybe IP
+  } deriving (Show, Eq, Typeable)
+
+extHandshakeId :: ExtendedMessageId
+extHandshakeId = 0
+
+-- | Default 'Request' queue size.
+defaultQueueLength :: Int
+defaultQueueLength = 1
+
+-- | All fields are empty.
+instance Default ExtendedHandshake where
+  def = ExtendedHandshake def def def def def def def def
+
+instance Monoid ExtendedHandshake where
+    mempty = def { ehsCaps = mempty }
+    mappend old new = ExtendedHandshake {
+      ehsCaps         = ehsCaps old <> ehsCaps new,
+      ehsIPv4         = ehsIPv4 old         `mergeOld` ehsIPv4 new,
+      ehsIPv6         = ehsIPv6 old         `mergeOld` ehsIPv6 new,
+      ehsMetadataSize = ehsMetadataSize old `mergeNew` ehsMetadataSize new,
+      ehsPort         = ehsPort old         `mergeOld` ehsPort new,
+      ehsQueueLength  = ehsQueueLength old  `mergeNew` ehsQueueLength new,
+      ehsVersion      = ehsVersion old      `mergeOld` ehsVersion new,
+      ehsYourIp       = ehsYourIp old       `mergeOld` ehsYourIp new
+    }
+      where
+        mergeOld mold mnew = mold <|> mnew
+        mergeNew mold mnew = mnew <|> mold
+
+
+instance BEncode ExtendedHandshake where
+  toBEncode ExtendedHandshake {..} = toDict $
+       "ipv4"   .=? (S.encode <$> ehsIPv4)
+    .: "ipv6"   .=? (S.encode <$> ehsIPv6)
+    .: "m"      .=! ehsCaps
+    .: "metadata_size" .=? ehsMetadataSize
+    .: "p"      .=? ehsPort
+    .: "reqq"   .=? ehsQueueLength
+    .: "v"      .=? ehsVersion
+    .: "yourip" .=? (runPut <$> either put put <$> toEither <$> ehsYourIp)
+    .: endDict
+    where
+      toEither (IPv4 v4) = Left v4
+      toEither (IPv6 v6) = Right v6
+
+  fromBEncode = fromDict $ ExtendedHandshake
+    <$>? "ipv4"
+    <*>? "ipv6"
+    <*>! "m"
+    <*>? "metadata_size"
+    <*>? "p"
+    <*>? "reqq"
+    <*>? "v"
+    <*> (opt "yourip" >>= getYourIp)
+
+getYourIp :: Maybe BValue -> BE.Get (Maybe IP)
+getYourIp f =
+  return $ do
+    BString ip <- f
+    either (const Nothing) Just $
+           case BS.length ip of
+             4  -> IPv4 <$> S.decode ip
+             16 -> IPv6 <$> S.decode ip
+             _ -> fail ""
+
+instance Pretty ExtendedHandshake where
+  pPrint = PP.text . show
+
+-- | NOTE: Approximated 'stats'.
+instance PeerMessage ExtendedHandshake where
+  envelop  c = envelop c . EHandshake
+  {-# INLINE envelop #-}
+
+  requires _ = Just ExtExtended
+  {-# INLINE requires #-}
+
+  stats _ = ByteStats (4 + 1 + 1) 100 {- is it ok? -} 0 -- FIXME
+  {-# INLINE stats #-}
+
+-- | Set default values and the specified 'ExtendedCaps'.
+nullExtendedHandshake :: ExtendedCaps -> ExtendedHandshake
+nullExtendedHandshake caps = ExtendedHandshake
+    { ehsIPv4         = Nothing
+    , ehsIPv6         = Nothing
+    , ehsCaps         = caps
+    , ehsMetadataSize = Nothing
+    , ehsPort         = Nothing
+    , ehsQueueLength  = Just defaultQueueLength
+    , ehsVersion      = Just $ T.pack $ render $ pPrint libFingerprint
+    , ehsYourIp       = Nothing
+    }
+
+{-----------------------------------------------------------------------
+-- Metadata exchange extension
+-----------------------------------------------------------------------}
+
+-- | A peer MUST verify that any piece it sends passes the info-hash
+-- verification. i.e. until the peer has the entire metadata, it
+-- cannot run SHA-1 to verify that it yields the same hash as the
+-- info-hash.
+--
+data ExtendedMetadata
+    -- | This message requests the a specified metadata piece. The
+    -- response to this message, from a peer supporting the extension,
+    -- is either a 'MetadataReject' or a 'MetadataData' message.
+  = MetadataRequest PieceIx
+
+    -- | If sender requested a valid 'PieceIx' and receiver have the
+    -- corresponding piece then receiver should respond with this
+    -- message.
+  | MetadataData
+    { -- | A piece of 'Data.Torrent.InfoDict'.
+      piece     :: P.Piece BS.ByteString
+
+      -- | This key has the same semantics as the 'ehsMetadataSize' in
+      -- the 'ExtendedHandshake' — it is size of the torrent info
+      -- dict.
+    , totalSize :: Int
+    }
+
+    -- | Peers that do not have the entire metadata MUST respond with
+    -- a reject message to any metadata request.
+    --
+    --   Clients MAY implement flood protection by rejecting request
+    --   messages after a certain number of them have been
+    --   served. Typically the number of pieces of metadata times a
+    --   factor.
+  | MetadataReject  PieceIx
+
+    -- | Reserved. By specification we should ignore unknown metadata
+    -- messages.
+  | MetadataUnknown BValue
+    deriving (Show, Eq, Typeable)
+
+-- | Extended metadata message id used in 'msg_type_key'.
+type MetadataId = Int
+
+msg_type_key, piece_key, total_size_key :: BKey
+msg_type_key   = "msg_type"
+piece_key      = "piece"
+total_size_key = "total_size"
+
+-- | BEP9 compatible encoding.
+instance BEncode ExtendedMetadata where
+  toBEncode (MetadataRequest pix) = toDict $
+       msg_type_key   .=! (0 :: MetadataId)
+    .: piece_key      .=! pix
+    .: endDict
+  toBEncode (MetadataData (P.Piece pix _) totalSize) = toDict $
+       msg_type_key   .=! (1 :: MetadataId)
+    .: piece_key      .=! pix
+    .: total_size_key .=! totalSize
+    .: endDict
+  toBEncode (MetadataReject  pix)  = toDict $
+       msg_type_key   .=! (2 :: MetadataId)
+    .: piece_key      .=! pix
+    .: endDict
+  toBEncode (MetadataUnknown bval) = bval
+
+  fromBEncode bval = (`fromDict` bval) $ do
+    mid <- field $ req msg_type_key
+    case mid :: MetadataId of
+      0 -> MetadataRequest <$>! piece_key
+      1 -> metadataData    <$>! piece_key <*>! total_size_key
+      2 -> MetadataReject  <$>! piece_key
+      _ -> pure (MetadataUnknown bval)
+   where
+     metadataData pix s = MetadataData (P.Piece pix BS.empty) s
+
+-- | Piece data bytes are omitted.
+instance Pretty ExtendedMetadata where
+  pPrint (MetadataRequest pix  ) = "Request" <+> PP.int    pix
+  pPrint (MetadataData    p   t) = "Data"    <+>    pPrint p   <+> PP.int t
+  pPrint (MetadataReject  pix  ) = "Reject"  <+> PP.int    pix
+  pPrint (MetadataUnknown bval ) = "Unknown" <+> ppBEncode bval
+
+-- | NOTE: Approximated 'stats'.
+instance PeerMessage ExtendedMetadata where
+  envelop  c = envelop c . EMetadata (remoteMessageId ExtMetadata c)
+  {-# INLINE envelop #-}
+
+  requires _ = Just ExtExtended
+  {-# INLINE requires #-}
+
+  stats (MetadataRequest _) = ByteStats (4 + 1 + 1) {- ~ -} 25 0
+  stats (MetadataData p  _) = ByteStats (4 + 1 + 1) {- ~ -} 41 $
+                              BS.length (P.pieceData p)
+  stats (MetadataReject  _) = ByteStats (4 + 1 + 1) {- ~ -} 25 0
+  stats (MetadataUnknown _) = ByteStats (4 + 1 + 1) {- ? -} 0  0
+
+-- | All 'Piece's in 'MetadataData' messages MUST have size equal to
+-- this value. The last trailing piece can be shorter.
+metadataPieceSize :: PieceSize
+metadataPieceSize = 16 * 1024
+
+isLastPiece :: P.Piece a -> Int -> Bool
+isLastPiece P.Piece {..} total = succ pieceIndex == pcnt
+  where
+    pcnt = q + if r > 0 then 1 else 0
+    (q, r) = quotRem total metadataPieceSize
+
+-- TODO we can check if the piece payload bytestring have appropriate
+-- length; otherwise serialization MUST fail.
+isValidPiece :: P.Piece BL.ByteString -> Int -> Bool
+isValidPiece p @ P.Piece {..} total
+  | isLastPiece p total = pieceSize p <= metadataPieceSize
+  |       otherwise     = pieceSize p == metadataPieceSize
+
+setMetadataPayload :: BS.ByteString -> ExtendedMetadata -> ExtendedMetadata
+setMetadataPayload bs (MetadataData (P.Piece pix _) t) =
+  MetadataData (P.Piece pix bs) t
+setMetadataPayload _   msg               = msg
+
+getMetadataPayload :: ExtendedMetadata -> Maybe BS.ByteString
+getMetadataPayload (MetadataData (P.Piece _ bs) _) = Just bs
+getMetadataPayload _                                  = Nothing
+
+-- | Metadata BDict usually contain only 'msg_type_key', 'piece_key'
+-- and 'total_size_key' fields so it normally should take less than
+-- 100 bytes. This limit is two order of magnitude larger to be
+-- permissive to 'MetadataUnknown' messages.
+--
+--   See 'maxMessageSize' for further explanation.
+--
+maxMetadataBDictSize :: Int
+maxMetadataBDictSize = 16 * 1024
+
+maxMetadataSize :: Int
+maxMetadataSize = maxMetadataBDictSize + metadataPieceSize
+
+-- to make MetadataData constructor fields a little bit prettier we
+-- cheat here: first we read empty 'pieceData' from bdict, but then we
+-- fill that field with the actual piece data — trailing bytes of
+-- the message
+getMetadata :: Int -> S.Get ExtendedMetadata
+getMetadata len
+    | len > maxMetadataSize = fail $ parseError "size exceeded limit"
+    | otherwise = do
+      bs <- getByteString len
+      parseRes $ BS.parse BE.parser bs
+  where
+    parseError reason = "unable to parse metadata message: " ++ reason
+
+    parseRes (BS.Fail _ _ m) = fail $ parseError $ "bdict: " ++ m
+    parseRes (BS.Partial _)  = fail $ parseError   "bdict: not enough bytes"
+    parseRes (BS.Done piece bvalueBS)
+      | BS.length piece > metadataPieceSize
+      = fail "infodict piece: size exceeded limit"
+      | otherwise = do
+        metadata <- either (fail . parseError) pure $ fromBEncode bvalueBS
+        return $ setMetadataPayload piece metadata
+
+putMetadata :: ExtendedMetadata -> BL.ByteString
+putMetadata msg
+  | Just bs <- getMetadataPayload msg = BE.encode msg <> BL.fromStrict bs
+  |              otherwise            = BE.encode msg
+
+-- | Allows a requesting peer to send 2 'MetadataRequest's for the
+--   each piece.
+--
+--   See 'Network.BitTorrent.Wire.Options.metadataFactor' for
+--   explanation why do we need this limit.
+defaultMetadataFactor :: Int
+defaultMetadataFactor = 2
+
+-- | Usually torrent size do not exceed 1MB. This value limit torrent
+--   /content/ size to about 8TB.
+--
+--   See 'Network.BitTorrent.Wire.Options.maxInfoDictSize' for
+--   explanation why do we need this limit.
+defaultMaxInfoDictSize :: Int
+defaultMaxInfoDictSize = 10 * 1024 * 1024
+
+{-----------------------------------------------------------------------
+--  Extension protocol messages
+-----------------------------------------------------------------------}
+
+-- | For more info see <http://www.bittorrent.org/beps/bep_0010.html>
+data ExtendedMessage
+  = EHandshake ExtendedHandshake
+  | EMetadata  ExtendedMessageId ExtendedMetadata
+  | EUnknown   ExtendedMessageId BS.ByteString
+    deriving (Show, Eq, Typeable)
+
+instance Pretty ExtendedMessage where
+  pPrint (EHandshake     ehs) = pPrint ehs
+  pPrint (EMetadata  _   msg) = "Metadata" <+> pPrint msg
+  pPrint (EUnknown   mid _  ) = "Unknown"  <+> PP.text (show mid)
+
+instance PeerMessage ExtendedMessage where
+  envelop  _ = Extended
+  {-# INLINE envelop #-}
+
+  requires _ = Just ExtExtended
+  {-# INLINE requires #-}
+
+  stats (EHandshake  hs)  = stats hs
+  stats (EMetadata _ msg) = stats msg
+  stats (EUnknown  _ msg) = ByteStats (4 + 1 + 1) (BS.length msg) 0
+
+{-----------------------------------------------------------------------
+-- The message datatype
+-----------------------------------------------------------------------}
+
+type MessageId = Word8
+
+-- | Messages used in communication between peers.
+--
+--   Note: If some extensions are disabled (not present in extension
+--   mask) and client receive message used by the disabled
+--   extension then the client MUST close the connection.
+--
+data Message
+    -- | Peers may close the TCP connection if they have not received
+    -- any messages for a given period of time, generally 2
+    -- minutes. Thus, the KeepAlive message is sent to keep the
+    -- connection between two peers alive, if no /other/ message has
+    -- been sent in a given period of time.
+  = KeepAlive
+  | Status    !StatusUpdate -- ^ Messages used to update peer status.
+  | Available !Available    -- ^ Messages used to inform availability.
+  | Transfer  !Transfer     -- ^ Messages used to transfer 'Block's.
+
+    -- | Peer receiving a handshake indicating the remote peer
+    -- supports the 'ExtDHT' should send a 'Port' message. Peers that
+    -- receive this message should attempt to ping the node on the
+    -- received port and IP address of the remote peer.
+  | Port      !PortNumber
+  | Fast      !FastMessage
+  | Extended  !ExtendedMessage
+    deriving (Show, Eq)
+
+instance Default Message where
+  def = KeepAlive
+  {-# INLINE def #-}
+
+-- | Payload bytes are omitted.
+instance Pretty Message where
+  pPrint (KeepAlive  ) = "Keep alive"
+  pPrint (Status    m) = "Status" <+> pPrint m
+  pPrint (Available m) = pPrint m
+  pPrint (Transfer  m) = pPrint m
+  pPrint (Port      p) = "Port" <+> int (fromEnum p)
+  pPrint (Fast      m) = pPrint m
+  pPrint (Extended  m) = pPrint m
+
+instance PeerMessage Message where
+  envelop _ = id
+  {-# INLINE envelop #-}
+
+  requires  KeepAlive    = Nothing
+  requires (Status    _) = Nothing
+  requires (Available _) = Nothing
+  requires (Transfer  _) = Nothing
+  requires (Port      _) = Just ExtDHT
+  requires (Fast      _) = Just ExtFast
+  requires (Extended  _) = Just ExtExtended
+
+  stats  KeepAlive    = ByteStats 4 0 0
+  stats (Status    m) = stats m
+  stats (Available m) = stats m
+  stats (Transfer  m) = stats m
+  stats (Port      _) = ByteStats 5 2 0
+  stats (Fast      m) = stats m
+  stats (Extended  m) = stats m
+
+-- | PORT message.
+instance PeerMessage PortNumber where
+  envelop  _ = Port
+  {-# INLINE envelop #-}
+
+  requires _ = Just ExtDHT
+  {-# INLINE requires #-}
+
+-- | How long /this/ peer should wait before dropping connection, in
+-- seconds.
+defaultKeepAliveTimeout :: Int
+defaultKeepAliveTimeout = 2 * 60
+
+-- | How often /this/ peer should send 'KeepAlive' messages, in
+-- seconds.
+defaultKeepAliveInterval :: Int
+defaultKeepAliveInterval = 60
+
+getInt :: S.Get Int
+getInt = fromIntegral <$> S.getWord32be
+{-# INLINE getInt #-}
+
+putInt :: S.Putter Int
+putInt = S.putWord32be . fromIntegral
+{-# INLINE putInt #-}
+
+-- | This limit should protect against "out-of-memory" attacks: if a
+-- malicious peer have sent a long varlength message then receiver can
+-- accumulate too long bytestring in the 'Get'.
+--
+--   Normal messages should never exceed this limits.
+--
+--   See also 'maxBitfieldSize', 'maxBlockSize' limits.
+--
+maxMessageSize :: Int
+maxMessageSize = 20 + 1024 * 1024
+
+-- | This also limit max torrent size to:
+--
+--     max_bitfield_size * piece_ix_per_byte * max_piece_size =
+--     2 ^ 20 * 8 * 1MB =
+--     8TB
+--
+maxBitfieldSize :: Int
+maxBitfieldSize = 1024 * 1024
+
+getBitfield :: Int -> S.Get Bitfield
+getBitfield len
+  | len > maxBitfieldSize = fail "BITFIELD message size exceeded limit"
+  |       otherwise       = fromBitmap <$> getByteString len
+
+maxBlockSize :: Int
+maxBlockSize = 4 * defaultTransferSize
+
+getBlock :: Int -> S.Get (Block BL.ByteString)
+getBlock len
+  | len > maxBlockSize = fail "BLOCK message size exceeded limit"
+  | otherwise = Block <$> getInt <*> getInt
+                      <*> getLazyByteString (fromIntegral len)
+{-# INLINE getBlock #-}
+
+instance Serialize Message where
+  get = do
+    len <- getInt
+
+    when (len > maxMessageSize) $ do
+      fail "message body size exceeded the limit"
+
+    if len == 0 then return KeepAlive
+      else do
+        mid <- S.getWord8
+        case mid of
+          0x00 -> return $ Status (Choking True)
+          0x01 -> return $ Status (Choking False)
+          0x02 -> return $ Status (Interested True)
+          0x03 -> return $ Status (Interested False)
+          0x04 -> (Available . Have)     <$> getInt
+          0x05 -> (Available . Bitfield) <$> getBitfield (pred len)
+          0x06 -> (Transfer  . Request)  <$> S.get
+          0x07 -> (Transfer  . Piece)    <$> getBlock (len - 9)
+          0x08 -> (Transfer  . Cancel)   <$> S.get
+          0x09 -> Port <$> S.get
+          0x0D -> (Fast . SuggestPiece) <$> getInt
+          0x0E -> return $ Fast HaveAll
+          0x0F -> return $ Fast HaveNone
+          0x10 -> (Fast . RejectRequest) <$> S.get
+          0x11 -> (Fast . AllowedFast)   <$> getInt
+          0x14 -> Extended <$> getExtendedMessage (pred len)
+          _    -> do
+            rm <- S.remaining >>= S.getBytes
+            fail $ "unknown message ID: " ++ show mid ++ "\n"
+                ++ "remaining available bytes: " ++ show rm
+
+  put  KeepAlive      = putInt 0
+  put (Status    msg) = putStatus    msg
+  put (Available msg) = putAvailable msg
+  put (Transfer  msg) = putTransfer  msg
+  put (Port      p  ) = putPort      p
+  put (Fast      msg) = putFast      msg
+  put (Extended  m )  = putExtendedMessage m
+
+statusUpdateId :: StatusUpdate -> MessageId
+statusUpdateId (Choking    choking) = fromIntegral (0 + fromEnum choking)
+statusUpdateId (Interested choking) = fromIntegral (2 + fromEnum choking)
+
+putStatus :: Putter StatusUpdate
+putStatus su = do
+  putInt 1
+  putWord8 (statusUpdateId su)
+
+putAvailable :: Putter Available
+putAvailable (Have i)      = do
+  putInt 5
+  putWord8 0x04
+  putInt i
+putAvailable (Bitfield (toBitmap -> bs)) = do
+  putInt $ 1 + fromIntegral (BL.length bs)
+  putWord8 0x05
+  putLazyByteString bs
+
+putBlock :: Putter (Block BL.ByteString)
+putBlock Block {..} = do
+  putInt blkPiece
+  putInt blkOffset
+  putLazyByteString blkData
+
+putTransfer :: Putter Transfer
+putTransfer (Request blk) = putInt 13 >> S.putWord8 0x06 >> S.put blk
+putTransfer (Piece   blk) = do
+  putInt (9 + blockSize blk)
+  putWord8 0x07
+  putBlock blk
+putTransfer (Cancel  blk)  = putInt 13 >> S.putWord8 0x08 >> S.put blk
+
+putPort :: Putter PortNumber
+putPort p = do
+  putInt 3
+  putWord8 0x09
+  put p
+
+putFast :: Putter FastMessage
+putFast  HaveAll           = putInt 1  >> putWord8 0x0E
+putFast  HaveNone          = putInt 1  >> putWord8 0x0F
+putFast (SuggestPiece pix) = putInt 5  >> putWord8 0x0D >> putInt pix
+putFast (RejectRequest i ) = putInt 13 >> putWord8 0x10 >> put i
+putFast (AllowedFast   i ) = putInt 5  >> putWord8 0x11 >> putInt i
+
+maxEHandshakeSize :: Int
+maxEHandshakeSize = 16 * 1024
+
+getExtendedHandshake :: Int -> S.Get ExtendedHandshake
+getExtendedHandshake messageSize
+  | messageSize > maxEHandshakeSize
+  = fail "extended handshake size exceeded limit"
+  | otherwise = do
+    bs <- getByteString messageSize
+    either fail pure $ BE.decode bs
+
+maxEUnknownSize :: Int
+maxEUnknownSize = 64 * 1024
+
+getExtendedUnknown :: Int -> S.Get BS.ByteString
+getExtendedUnknown len
+  | len > maxEUnknownSize = fail "unknown extended message size exceeded limit"
+  |        otherwise      = getByteString len
+
+getExtendedMessage :: Int -> S.Get ExtendedMessage
+getExtendedMessage messageSize = do
+  msgId   <- getWord8
+  let msgBodySize = messageSize - 1
+  case msgId of
+    0 -> EHandshake      <$> getExtendedHandshake msgBodySize
+    1 -> EMetadata msgId <$> getMetadata          msgBodySize
+    _ -> EUnknown  msgId <$> getExtendedUnknown   msgBodySize
+
+-- | By spec.
+extendedMessageId :: MessageId
+extendedMessageId = 20
+
+putExt :: ExtendedMessageId -> BL.ByteString -> Put
+putExt mid lbs = do
+  putWord32be $ fromIntegral (1 + 1 + BL.length lbs)
+  putWord8 extendedMessageId
+  putWord8 mid
+  putLazyByteString lbs
+
+-- NOTE: in contrast to getExtendedMessage this function put length
+-- and message id too!
+putExtendedMessage :: Putter ExtendedMessage
+putExtendedMessage (EHandshake hs)     = putExt extHandshakeId $ BE.encode hs
+putExtendedMessage (EMetadata mid msg) = putExt mid $ putMetadata msg
+putExtendedMessage (EUnknown  mid  bs) = putExt mid $ BL.fromStrict bs