From 7641ea42ed7c35d9babfe270e5a93dd8cb4922ae Mon Sep 17 00:00:00 2001 From: Sam T Date: Sun, 30 Jun 2013 21:24:06 +0400 Subject: + Add instances for Binary. --- bittorrent.cabal | 1 + src/Network/BitTorrent/Exchange/Protocol.hs | 162 +++++++++++++++++++++------- 2 files changed, 123 insertions(+), 40 deletions(-) diff --git a/bittorrent.cabal b/bittorrent.cabal index 013133d5..02130385 100644 --- a/bittorrent.cabal +++ b/bittorrent.cabal @@ -89,6 +89,7 @@ library -- Encoding/Serialization , bencoding >= 0.1.0.1 , cereal >= 0.3 + , binary >= 0.5 , urlencoded >= 0.4 -- Time diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs index 4cf4685d..83774f06 100644 --- a/src/Network/BitTorrent/Exchange/Protocol.hs +++ b/src/Network/BitTorrent/Exchange/Protocol.hs @@ -72,9 +72,14 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as BC import qualified Data.ByteString.Lazy as Lazy import Data.Default -import Data.Serialize as S import Data.Int import Data.Word + +import Data.Binary as B +import Data.Binary.Get as B +import Data.Binary.Put as B +import Data.Serialize as S + import Text.PrettyPrint import Network @@ -117,18 +122,18 @@ data Handshake = Handshake { 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) + S.putWord8 (fromIntegral (B.length (hsProtocol hs))) + S.putByteString (hsProtocol hs) + S.putWord64be (hsReserved hs) + S.put (hsInfoHash hs) + S.put (hsPeerID hs) get = do - len <- getWord8 - Handshake <$> getBytes (fromIntegral len) - <*> getWord64be - <*> get - <*> get + len <- S.getWord8 + Handshake <$> S.getBytes (fromIntegral len) + <*> S.getWord64be + <*> S.get + <*> S.get -- | Extract capabilities from a peer handshake message. handshakeCaps :: Handshake -> Capabilities @@ -204,14 +209,22 @@ data BlockIx = BlockIx { , ixLength :: {-# UNPACK #-} !Int } deriving (Show, Eq) -getInt :: Get Int -getInt = fromIntegral <$> getWord32be +getInt :: S.Get Int +getInt = fromIntegral <$> S.getWord32be {-# INLINE getInt #-} -putInt :: Putter Int -putInt = putWord32be . fromIntegral +putInt :: S.Putter Int +putInt = S.putWord32be . fromIntegral {-# INLINE putInt #-} +getIntB :: B.Get Int +getIntB = fromIntegral <$> B.getWord32be +{-# INLINE getIntB #-} + +putIntB :: Int -> B.Put +putIntB = B.putWord32be . fromIntegral +{-# INLINE putIntB #-} + instance Serialize BlockIx where {-# SPECIALIZE instance Serialize BlockIx #-} get = BlockIx <$> getInt <*> getInt <*> getInt @@ -222,6 +235,16 @@ instance Serialize BlockIx where putInt (ixLength i) {-# INLINE put #-} +instance Binary BlockIx where + {-# SPECIALIZE instance Binary BlockIx #-} + get = BlockIx <$> getIntB <*> getIntB <*> getIntB + {-# INLINE get #-} + + put BlockIx {..} = do + putIntB ixPiece + putIntB ixOffset + putIntB ixLength + -- | Format block index in human readable form. ppBlockIx :: BlockIx -> Doc ppBlockIx BlockIx {..} = @@ -355,61 +378,120 @@ instance Serialize Message where -- _ <- lookAhead $ ensure len if len == 0 then return KeepAlive else do - mid <- getWord8 + mid <- S.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 + 0x05 -> (Bitfield . fromBitmap) <$> S.getByteString (pred len) + 0x06 -> Request <$> S.get 0x07 -> Piece <$> getBlock (len - 9) - 0x08 -> Cancel <$> get - 0x09 -> (Port . fromIntegral) <$> getWord16be + 0x08 -> Cancel <$> S.get + 0x09 -> (Port . fromIntegral) <$> S.getWord16be 0x0E -> return HaveAll 0x0F -> return HaveNone 0x0D -> SuggestPiece <$> getInt - 0x10 -> RejectRequest <$> get + 0x10 -> RejectRequest <$> S.get 0x11 -> AllowedFast <$> getInt _ -> do - rm <- remaining >>= getBytes + rm <- S.remaining >>= S.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 + getBlock :: Int -> S.Get Block + getBlock len = Block <$> getInt <*> getInt <*> S.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 + put Choke = putInt 1 >> S.putWord8 0x00 + put Unchoke = putInt 1 >> S.putWord8 0x01 + put Interested = putInt 1 >> S.putWord8 0x02 + put NotInterested = putInt 1 >> S.putWord8 0x03 + put (Have i) = putInt 5 >> S.putWord8 0x04 >> putInt i + put (Bitfield bf) = putInt l >> S.putWord8 0x05 >> S.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 + put (Request blk) = putInt 13 >> S.putWord8 0x06 >> S.put blk + put (Piece blk) = putInt l >> S.putWord8 0x07 >> putBlock where l = 9 + B.length (blkData blk) {-# INLINE l #-} putBlock = do putInt (blkPiece blk) putInt (blkOffset blk) - putByteString (blkData blk) + S.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 i ) = putInt 13 >> putWord8 0x10 >> put i - put (AllowedFast i ) = putInt 5 >> putWord8 0x11 >> putInt i + put (Cancel blk) = putInt 13 >> S.putWord8 0x08 >> S.put blk + put (Port p ) = putInt 3 >> S.putWord8 0x09 >> S.putWord16be (fromIntegral p) + put HaveAll = putInt 1 >> S.putWord8 0x0E + put HaveNone = putInt 1 >> S.putWord8 0x0F + put (SuggestPiece pix) = putInt 5 >> S.putWord8 0x0D >> putInt pix + put (RejectRequest i ) = putInt 13 >> S.putWord8 0x10 >> S.put i + put (AllowedFast i ) = putInt 5 >> S.putWord8 0x11 >> putInt i + +instance Binary Message where + get = do + len <- undefined --getInt +-- _ <- lookAhead $ ensure len + if len == 0 then return KeepAlive + else do + mid <- B.getWord8 + case mid of + 0x00 -> return Choke + 0x01 -> return Unchoke + 0x02 -> return Interested + 0x03 -> return NotInterested + 0x04 -> Have <$> getIntB + 0x05 -> (Bitfield . fromBitmap) <$> B.getByteString (pred len) + 0x06 -> Request <$> B.get + 0x07 -> Piece <$> getBlock (len - 9) + 0x08 -> Cancel <$> B.get + 0x09 -> (Port . fromIntegral) <$> B.getWord16be + 0x0E -> return HaveAll + 0x0F -> return HaveNone + 0x0D -> SuggestPiece <$> getIntB + 0x10 -> RejectRequest <$> B.get + 0x11 -> AllowedFast <$> getIntB + _ -> do + rm <- B.remaining >>= B.getBytes . fromIntegral + fail $ "unknown message ID: " ++ show mid ++ "\n" + ++ "remaining available bytes: " ++ show rm + + where + getBlock :: Int -> B.Get Block + getBlock len = Block <$> getIntB <*> getIntB <*> B.getBytes len + {-# INLINE getBlock #-} + + put KeepAlive = putIntB 0 + put Choke = putIntB 1 >> B.putWord8 0x00 + put Unchoke = putIntB 1 >> B.putWord8 0x01 + put Interested = putIntB 1 >> B.putWord8 0x02 + put NotInterested = putIntB 1 >> B.putWord8 0x03 + put (Have i) = putIntB 5 >> B.putWord8 0x04 >> putIntB i + put (Bitfield bf) = putIntB l >> B.putWord8 0x05 >> B.putLazyByteString b + where b = toBitmap bf + l = succ (fromIntegral (Lazy.length b)) + {-# INLINE l #-} + put (Request blk) = putIntB 13 >> B.putWord8 0x06 >> B.put blk + put (Piece blk) = putIntB l >> B.putWord8 0x07 >> putBlock + where l = 9 + B.length (blkData blk) + {-# INLINE l #-} + putBlock = do putIntB (blkPiece blk) + putIntB (blkOffset blk) + B.putByteString (blkData blk) + {-# INLINE putBlock #-} + put (Cancel blk) = putIntB 13 >> B.putWord8 0x08 >> B.put blk + put (Port p ) = putIntB 3 >> B.putWord8 0x09 >> B.putWord16be (fromIntegral p) + put HaveAll = putIntB 1 >> B.putWord8 0x0E + put HaveNone = putIntB 1 >> B.putWord8 0x0F + put (SuggestPiece pix) = putIntB 5 >> B.putWord8 0x0D >> putIntB pix + put (RejectRequest i ) = putIntB 13 >> B.putWord8 0x10 >> B.put i + put (AllowedFast i ) = putIntB 5 >> B.putWord8 0x11 >> putIntB i -- | Format messages in human readable form. Note that output is -- compact and suitable for logging: only useful information but not -- cgit v1.2.3