From 1e8a6a7d5267811d035afda764e90092eb0e994c Mon Sep 17 00:00:00 2001 From: Sam Truzjan Date: Thu, 5 Dec 2013 03:22:00 +0400 Subject: Add BEP9 messages --- src/Network/BitTorrent/Exchange/Message.hs | 193 +++++++++++++++++++++-------- 1 file changed, 142 insertions(+), 51 deletions(-) (limited to 'src') diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs index 0a535517..2f85d729 100644 --- a/src/Network/BitTorrent/Exchange/Message.hs +++ b/src/Network/BitTorrent/Exchange/Message.hs @@ -25,48 +25,50 @@ -- For more infomation see: -- -- -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS -fno-warn-orphans #-} module Network.BitTorrent.Exchange.Message - ( -- * Extensions + ( -- * Capabilities Extension (..) , Caps - , requires - , allowed , toCaps , fromCaps + , allowed -- * Handshake , Handshake(..) , defaultHandshake , defaultBTProtocol + , handshakeSize , handshakeMaxSize - -- * TODO remove this section from this module - , handshake - , recvHandshake - , sendHandshake - -- * Messages , Message (..) + , PeerMessage (..) + , requires + + -- ** Core messages , StatusUpdate (..) , RegularMessage (..) - -- * Fast extension + -- ** Fast extension , FastMessage (..) - -- * Extension protocol + -- ** Extension protocol + , ExtendedMessage (..) , ExtendedExtension , ExtendedCaps (..) , ExtendedHandshake (..) - , ExtendedMessage (..) + , ExtendedMetadata (..) ) where import Control.Applicative -import Control.Exception -import Control.Monad import Data.BEncode as BE +import Data.BEncode.BDict as BE +import Data.BEncode.Internal (ppBEncode) import Data.Bits import Data.ByteString as BS import Data.ByteString.Char8 as BC @@ -75,9 +77,11 @@ import Data.Default import Data.IntMap as IM import Data.List as L import Data.Monoid +import Data.Ord import Data.Serialize as S import Data.Text as T import Data.Typeable +import Data.Tuple import Data.Word import Network import Network.Socket hiding (KeepAlive) @@ -219,36 +223,13 @@ defaultBTProtocol = "BitTorrent protocol" defaultHandshake :: InfoHash -> PeerId -> Handshake defaultHandshake = Handshake defaultBTProtocol def --- | TODO remove socket stuff to corresponding module -sendHandshake :: Socket -> Handshake -> IO () -sendHandshake sock hs = sendAll sock (S.encode hs) - -recvHandshake :: Socket -> IO Handshake -recvHandshake sock = do - header <- BS.recv sock 1 - unless (BS.length header == 1) $ - throw $ userError "Unable to receive handshake header." - - let protocolLen = BS.head header - let restLen = handshakeSize protocolLen - 1 - - body <- BS.recv sock restLen - let resp = BS.cons protocolLen body - either (throwIO . userError) return $ S.decode resp - --- | Handshaking with a peer specified by the second argument. -handshake :: Socket -> Handshake -> IO Handshake -handshake sock hs = do - sendHandshake sock hs - hs' <- recvHandshake sock - when (hsInfoHash hs /= hsInfoHash hs') $ do - throwIO $ userError "Handshake info hash do not match." - return hs' - {----------------------------------------------------------------------- Regular messages -----------------------------------------------------------------------} +class PeerMessage a where + envelop :: ExtendedCaps -> a -> Message + data StatusUpdate = Choke | Unchoke @@ -259,6 +240,9 @@ data StatusUpdate instance Pretty StatusUpdate where pretty = text . show +instance PeerMessage StatusUpdate where + envelop _ = Status + data RegularMessage = -- | Zero-based index of a piece that has just been successfully -- downloaded and verified via the hash. @@ -298,6 +282,18 @@ instance Pretty RegularMessage where pretty (Piece blk) = "Piece" <+> pretty blk pretty (Cancel i ) = "Cancel" <+> pretty i +instance PeerMessage RegularMessage where + envelop _ = Regular + +instance PeerMessage Bitfield where + envelop c = envelop c . Bitfield + +instance PeerMessage BlockIx where + envelop c = envelop c . Request + +instance PeerMessage (Block BL.ByteString) where + envelop c = envelop c . Piece + -- | BEP6 messages. data FastMessage = -- | If a peer have all pieces it might send the 'HaveAll' message @@ -329,26 +325,58 @@ instance Pretty FastMessage where pretty (RejectRequest bix) = "Reject" <+> pretty bix pretty (AllowedFast pix) = "Allowed fast" <+> int pix +instance PeerMessage FastMessage where + envelop _ = Fast + {----------------------------------------------------------------------- -- Extended messages -----------------------------------------------------------------------} -type ExtendedExtension = () - type ExtendedMessageId = Word8 - type ExtendedIdMap = IntMap --- | The extension IDs must be stored for every peer, becuase every +data ExtendedExtension + = ExtMetadata -- ^ BEP 9 + deriving (Show, Eq, Typeable) + +instance Pretty ExtendedExtension where + pretty ExtMetadata = "Extension for Peers to Send Metadata Files" + +extId :: ExtendedExtension -> ExtendedMessageId +extId ExtMetadata = 1 +{-# INLINE extId #-} + +extString :: ExtendedExtension -> BS.ByteString +extString ExtMetadata = "ut_metadata" +{-# INLINE extString #-} + +fromS :: BS.ByteString -> ExtendedExtension +fromS "ut_metadata" = ExtMetadata + +-- | The extension IDs must be stored for every peer, because every -- peer may have different IDs for the same extension. -- newtype ExtendedCaps = ExtendedCaps { extendedCaps :: ExtendedIdMap ExtendedExtension - } deriving (Show, Eq) + } deriving (Show, Eq, Monoid) + +-- | Empty set. +instance Default ExtendedCaps where + def = ExtendedCaps IM.empty + +instance Pretty ExtendedCaps where + pretty = ppBEncode . toBEncode instance BEncode ExtendedCaps where - fromBEncode = undefined - toBEncode = undefined + toBEncode = BDict . BE.fromAscList . L.sortBy (comparing fst) + . L.map mkPair . IM.toList . extendedCaps + where + mkPair (eid, ex) = (extString ex, toBEncode eid) + + fromBEncode (BDict bd) = ExtendedCaps <$> undefined + + fromBEncode _ = decodingError "ExtendedCaps" + -- | This message should be sent immediately after the standard -- bittorrent handshake to any peer that supports this extension @@ -385,10 +413,13 @@ data ExtendedHandshake = ExtendedHandshake -- , yourip :: Maybe (Either HostAddress HostAddress6) } deriving (Show, Eq, Typeable) +instance Default ExtendedHandshake where + def = ExtendedHandshake Nothing Nothing def Nothing Nothing Nothing + instance BEncode ExtendedHandshake where toBEncode ExtendedHandshake {..} = toDict $ - "ipv4" .=? ehsIPv4 - .: "ipv6" .=? ehsIPv6 + "ipv4" .=? ehsIPv4 -- FIXME invalid encoding + .: "ipv6" .=? ehsIPv6 -- FIXME invalid encoding .: "m" .=! ehsCaps .: "p" .=? ehsPort .: "reqq" .=? ehsQueueLength @@ -408,16 +439,67 @@ instance BEncode ExtendedHandshake where instance Pretty ExtendedHandshake where pretty = PP.text . show +instance PeerMessage ExtendedHandshake where + envelop c = envelop c . EHandshake + +{----------------------------------------------------------------------- +-- Metadata exchange +-----------------------------------------------------------------------} + +type MetadataId = Int + +pieceSize :: Int +pieceSize = 16 * 1024 + +data ExtendedMetadata + = MetadataRequest PieceIx + | MetadataData PieceIx Int + | MetadataReject PieceIx + | MetadataUnknown BValue + deriving (Show, Eq, Typeable) + +instance BEncode ExtendedMetadata where + toBEncode (MetadataRequest pix) = toDict $ + "msg_type" .=! (0 :: MetadataId) + .: "piece" .=! pix + .: endDict + toBEncode (MetadataData pix totalSize) = toDict $ + "msg_type" .=! (1 :: MetadataId) + .: "piece" .=! pix + .: "total_size" .=! totalSize + .: endDict + toBEncode (MetadataReject pix) = toDict $ + "msg_type" .=! (2 :: MetadataId) + .: "piece" .=! pix + .: endDict + toBEncode (MetadataUnknown bval) = bval + + fromBEncode = undefined + +instance Pretty ExtendedMetadata where + pretty (MetadataRequest pix ) = "Request" <+> PP.int pix + pretty (MetadataData pix s) = "Data" <+> PP.int pix <+> PP.int s + pretty (MetadataReject pix ) = "Reject" <+> PP.int pix + pretty (MetadataUnknown bval ) = ppBEncode bval + +instance PeerMessage ExtendedMetadata where + envelop c = envelop c . EMetadata + -- | For more info see data ExtendedMessage = EHandshake ExtendedHandshake + | EMetadata ExtendedMetadata | EUnknown ExtendedMessageId BS.ByteString - deriving (Show, Eq) + deriving (Show, Eq, Typeable) instance Pretty ExtendedMessage where pretty (EHandshake ehs) = pretty ehs + pretty (EMetadata msg) = pretty msg pretty (EUnknown mid _) = "Unknown" <+> PP.text (show mid) +instance PeerMessage ExtendedMessage where + envelop _ = Extended + {----------------------------------------------------------------------- -- The message datatype -----------------------------------------------------------------------} @@ -455,6 +537,12 @@ instance Pretty Message where pretty (Fast m) = pretty m pretty (Extended m) = pretty m +instance PeerMessage Message where + envelop _ = id + +instance PeerMessage PortNumber where + envelop _ = Port + -- | Can be used to check if this message is allowed to send\/recv in -- current session. requires :: Message -> Maybe Extension @@ -555,6 +643,7 @@ getExtendedMessage messageSize = do let msgBodySize = messageSize - 1 case msgId of 0 -> EHandshake <$> getExtendedHandshake msgBodySize + 1 -> EMetadata <$> undefined _ -> EUnknown msgId <$> getByteString msgBodySize extendedMessageId :: MessageId @@ -565,7 +654,9 @@ extendedMessageId = 20 putExtendedMessage :: ExtendedMessage -> S.Put putExtendedMessage (EHandshake hs) = do putExtendedMessage $ EUnknown 0 $ BL.toStrict $ BE.encode hs - +putExtendedMessage (EMetadata msg) = do + putExtendedMessage $ EUnknown (extId ExtMetadata) + $ BL.toStrict $ BE.encode msg putExtendedMessage (EUnknown mid bs) = do putWord32be $ fromIntegral (4 + 1 + BS.length bs) putWord8 extendedMessageId -- cgit v1.2.3