-- | -- Copyright : (c) Sam T. 2013 -- License : MIT -- Maintainer : pxqr.sta@gmail.com -- Stability : experimental -- Portability : portable -- -- 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. -- {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} module Network.BitTorrent.PeerWire.Protocol ( -- * Inital handshake Handshake(..), ppHandshake , handshake , handshakeCaps -- ** Defaults , defaultHandshake, defaultBTProtocol, defaultReserved , handshakeMaxSize -- * Block , PieceIx, BlockLIx, PieceLIx , BlockIx(..), ppBlockIx , Block(..), ppBlock ,blockSize , pieceIx, blockIx , blockRange, ixRange, isPiece -- ** Defaults , defaultBlockSize -- * Regular messages , Message(..) , ppMessage ) where import Control.Applicative import Control.Monad import Control.Exception import Data.ByteString (ByteString) import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as BC import qualified Data.ByteString.Lazy as Lazy import Data.Serialize as S import Data.Int import Data.Word import Text.PrettyPrint import Network import Network.Socket.ByteString import Data.Bitfield import Data.Torrent import Network.BitTorrent.Extension import Network.BitTorrent.Peer {----------------------------------------------------------------------- Handshake -----------------------------------------------------------------------} data Handshake = Handshake { -- | Identifier of the protocol. hsProtocol :: ByteString -- | Reserved bytes used to specify supported BEP's. , hsReserved :: Capabilities -- | 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 hs = do putWord8 (fromIntegral (B.length (hsProtocol hs))) putByteString (hsProtocol hs) putWord64be (hsReserved hs) put (hsInfoHash hs) put (hsPeerID hs) get = do len <- getWord8 Handshake <$> getBytes (fromIntegral len) <*> getWord64be <*> get <*> get handshakeCaps :: Handshake -> Capabilities handshakeCaps = hsReserved -- | Format handshake in human readable form. ppHandshake :: Handshake -> Doc ppHandshake Handshake {..} = text (BC.unpack hsProtocol) <+> ppClientInfo (clientInfo 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 255 -- | Default protocol string "BitTorrent protocol" as is. defaultBTProtocol :: ByteString defaultBTProtocol = "BitTorrent protocol" -- | Default reserved word is 0. defaultReserved :: Word64 defaultReserved = 0 -- | Length of info hash and peer id is unchecked, so it /should/ be equal 20. defaultHandshake :: InfoHash -> PeerID -> Handshake defaultHandshake = Handshake defaultBTProtocol defaultReserved -- | Handshaking with a peer specified by the second argument. handshake :: Socket -> Handshake -> IO Handshake handshake sock hs = do sendAll sock (S.encode hs) header <- recv sock 1 when (B.length header == 0) $ throw $ userError "Unable to receive handshake." let protocolLen = B.head header let restLen = handshakeSize protocolLen - 1 body <- recv sock restLen let resp = B.cons protocolLen body case checkIH (S.decode resp) of Right hs' -> return hs' Left msg -> throw $ userError msg where checkIH (Right hs') | hsInfoHash hs /= hsInfoHash hs' = Left "Handshake info hash do not match." checkIH x = x {----------------------------------------------------------------------- Blocks -----------------------------------------------------------------------} type BlockLIx = Int type PieceLIx = Int data BlockIx = BlockIx { -- | Zero-based piece index. ixPiece :: {-# UNPACK #-} !PieceLIx -- | Zero-based byte offset within the piece. , ixOffset :: {-# UNPACK #-} !Int -- | Block size starting from offset. , ixLength :: {-# UNPACK #-} !Int } deriving (Show, Eq) getInt :: Get Int getInt = fromIntegral <$> getWord32be {-# INLINE getInt #-} putInt :: Putter Int putInt = putWord32be . fromIntegral {-# INLINE putInt #-} instance Serialize BlockIx where {-# SPECIALIZE instance Serialize BlockIx #-} get = BlockIx <$> getInt <*> getInt <*> getInt {-# INLINE get #-} put ix = do putInt (ixPiece ix) putInt (ixOffset ix) putInt (ixLength ix) {-# INLINE put #-} ppBlockIx :: BlockIx -> Doc ppBlockIx BlockIx {..} = "piece = " <> int ixPiece <> "," <+> "offset = " <> int ixOffset <> "," <+> "length = " <> int ixLength data Block = Block { -- | Zero-based piece index. blkPiece :: !PieceLIx -- | Zero-based byte offset within the piece. , blkOffset :: !Int -- | Payload. , blkData :: !ByteString } deriving (Show, Eq) ppBlock :: Block -> Doc ppBlock = ppBlockIx . blockIx blockSize :: Block -> Int blockSize blk = B.length (blkData blk) -- | Widely used semi-official block size. defaultBlockSize :: Int defaultBlockSize = 16 * 1024 isPiece :: Int -> Block -> Bool isPiece pieceSize (Block i offset bs) = offset == 0 && B.length bs == pieceSize && i >= 0 {-# INLINE isPiece #-} pieceIx :: Int -> Int -> BlockIx pieceIx i = BlockIx i 0 {-# INLINE pieceIx #-} blockIx :: Block -> BlockIx blockIx = BlockIx <$> blkPiece <*> blkOffset <*> B.length . blkData blockRange :: (Num a, Integral a) => Int -> Block -> (a, a) blockRange pieceSize blk = (offset, offset + len) where offset = fromIntegral pieceSize * fromIntegral (blkPiece blk) + fromIntegral (blkOffset blk) len = fromIntegral (B.length (blkData blk)) {-# INLINE blockRange #-} {-# SPECIALIZE blockRange :: Int -> Block -> (Int64, Int64) #-} ixRange :: (Num a, Integral a) => Int -> BlockIx -> (a, a) ixRange pieceSize ix = (offset, offset + len) where offset = fromIntegral pieceSize * fromIntegral (ixPiece ix) + fromIntegral (ixOffset ix) len = fromIntegral (ixLength ix) {-# INLINE ixRange #-} {-# SPECIALIZE ixRange :: Int -> BlockIx -> (Int64, Int64) #-} {----------------------------------------------------------------------- Handshake -----------------------------------------------------------------------} -- | 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 = KeepAlive | Choke | Unchoke | Interested | NotInterested -- | 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 -- | 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 for a request for a block. | Piece !Block -- | Used to cancel block requests. It is typically -- used during "End Game". | Cancel !BlockIx | Port !PortNumber -- | BEP 6: Then peer have all pieces it might send the -- 'HaveAll' message instead of 'Bitfield' -- message. Used to save bandwidth. | HaveAll -- | BEP 6: Then peer have no pieces it might send -- 'HaveNone' message intead of 'Bitfield' -- message. Used to save bandwidth. | HaveNone -- | BEP 6: 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 -- | BEP 6: Notifies a requesting peer that its request -- will not be satisfied. | RejectRequest !BlockIx -- | BEP 6: 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 Serialize Message where get = do len <- getInt -- _ <- lookAhead $ ensure len if len == 0 then return KeepAlive else do mid <- getWord8 case mid of 0x00 -> return Choke 0x01 -> return Unchoke 0x02 -> return Interested 0x03 -> return NotInterested 0x04 -> Have <$> getInt 0x05 -> (Bitfield . fromBitmap) <$> getByteString (pred len) 0x06 -> Request <$> get 0x07 -> Piece <$> getBlock (len - 9) 0x08 -> Cancel <$> get 0x09 -> (Port . fromIntegral) <$> getWord16be 0x0E -> return HaveAll 0x0F -> return HaveNone 0x0D -> SuggestPiece <$> getInt 0x10 -> RejectRequest <$> get 0x11 -> AllowedFast <$> getInt _ -> do rm <- remaining >>= getBytes fail $ "unknown message ID: " ++ show mid ++ "\n" ++ "remaining available bytes: " ++ show rm where getBlock :: Int -> Get Block getBlock len = Block <$> getInt <*> getInt <*> getBytes len {-# INLINE getBlock #-} put KeepAlive = putInt 0 put Choke = putInt 1 >> putWord8 0x00 put Unchoke = putInt 1 >> putWord8 0x01 put Interested = putInt 1 >> putWord8 0x02 put NotInterested = putInt 1 >> putWord8 0x03 put (Have i) = putInt 5 >> putWord8 0x04 >> putInt i put (Bitfield bf) = putInt l >> putWord8 0x05 >> putLazyByteString b where b = toBitmap bf l = succ (fromIntegral (Lazy.length b)) {-# INLINE l #-} put (Request blk) = putInt 13 >> putWord8 0x06 >> put blk put (Piece blk) = putInt l >> putWord8 0x07 >> putBlock where l = 9 + B.length (blkData blk) {-# INLINE l #-} putBlock = do putInt (blkPiece blk) putInt (blkOffset blk) putByteString (blkData blk) {-# INLINE putBlock #-} put (Cancel blk) = putInt 13 >> putWord8 0x08 >> put blk put (Port p ) = putInt 3 >> putWord8 0x09 >> putWord16be (fromIntegral p) put HaveAll = putInt 1 >> putWord8 0x0E put HaveNone = putInt 1 >> putWord8 0x0F put (SuggestPiece pix) = putInt 5 >> putWord8 0x0D >> putInt pix put (RejectRequest ix) = putInt 13 >> putWord8 0x10 >> put ix put (AllowedFast ix) = putInt 5 >> putWord8 0x11 >> putInt ix -- | Format messages in human readable form. Note that output is -- compact and suitable for logging: only useful information but not -- payload bytes. -- ppMessage :: Message -> Doc ppMessage (Bitfield _) = "Bitfield" ppMessage (Piece blk) = "Piece" <+> ppBlock blk ppMessage (Cancel ix) = "Cancel" <+> ppBlockIx ix ppMessage (SuggestPiece pix) = "Suggest" <+> int pix ppMessage (RejectRequest ix) = "Reject" <+> ppBlockIx ix ppMessage msg = text (show msg)