summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--network-bittorrent.cabal3
-rw-r--r--src/Network/Torrent/PeerWire.hs103
-rw-r--r--src/Network/Torrent/PeerWire/Block.hs17
-rw-r--r--src/Network/Torrent/PeerWire/Message.hs91
4 files changed, 115 insertions, 99 deletions
diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index d55b5b50..6fa115e6 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -28,9 +28,10 @@ library
28 , Network.Torrent.Tracker 28 , Network.Torrent.Tracker
29 , Network.Torrent.Tracker.Scrape 29 , Network.Torrent.Tracker.Scrape
30 , Network.Torrent.PeerWire 30 , Network.Torrent.PeerWire
31 , Network.Torrent.PeerWire.Block
32 , Network.Torrent.PeerWire.Message
31 , Network.Torrent.PeerWire.Handshake 33 , Network.Torrent.PeerWire.Handshake
32 34
33
34 other-modules: 35 other-modules:
35 36
36 37
diff --git a/src/Network/Torrent/PeerWire.hs b/src/Network/Torrent/PeerWire.hs
index f84b7016..887e56d5 100644
--- a/src/Network/Torrent/PeerWire.hs
+++ b/src/Network/Torrent/PeerWire.hs
@@ -7,104 +7,11 @@
7-- 7--
8{-# LANGUAGE DoAndIfThenElse #-} 8{-# LANGUAGE DoAndIfThenElse #-}
9module Network.Torrent.PeerWire 9module Network.Torrent.PeerWire
10 ( module Network.Torrent.PeerWire.Handshake 10 ( module Network.Torrent.PeerWire.Block
11 , Message(..), Block(..), BlockIx(..), 11 , module Network.Torrent.PeerWire.Message
12 , module Network.Torrent.PeerWire.Handshake
12 ) where 13 ) where
13 14
15import Network.Torrent.PeerWire.Block
16import Network.Torrent.PeerWire.Message
14import Network.Torrent.PeerWire.Handshake 17import Network.Torrent.PeerWire.Handshake
15
16import Control.Applicative
17import Data.ByteString (ByteString)
18import qualified Data.ByteString as B
19
20import Data.Serialize
21
22data BlockIx = BlockIx {
23 ixPiece :: {-# UNPACK #-} !Int -- ^ Zero-based piece index.
24 , ixOffset :: {-# UNPACK #-} !Int -- ^ Zero-based byte offset within the piece.
25 , ixLength :: {-# UNPACK #-} !Int -- ^ Block size starting from offset.
26 } deriving (Show, Eq)
27
28data Block = Block {
29 blkPiece :: Int -- ^ Zero-based piece index.
30 , blkOffset :: Int -- ^ Zero-based byte offset within the piece.
31 , blkData :: ByteString -- ^ Payload.
32 } deriving (Show, Eq)
33
34-- TODO comment message constructors
35data Message = KeepAlive
36 | Choke
37 | Unchoke
38 | Interested
39 | NotInterested
40 | Have Int
41 | Bitfield ByteString
42 | Request BlockIx
43 | Piece Block
44 | Cancel BlockIx
45 | Port Int
46 deriving (Show, Eq)
47
48getInt :: Get Int
49getInt = fromIntegral <$> getWord32be
50{-# INLINE getInt #-}
51
52putInt :: Putter Int
53putInt = putWord32be . fromIntegral
54{-# INLINE putInt #-}
55
56instance Serialize BlockIx where
57 {-# SPECIALIZE instance Serialize BlockIx #-}
58 get = BlockIx <$> getInt <*> getInt <*> getInt
59 {-# INLINE get #-}
60
61 put ix = do putInt (ixPiece ix)
62 putInt (ixOffset ix)
63 putInt (ixLength ix)
64 {-# INLINE put #-}
65
66instance Serialize Message where
67 get = do
68 len <- getInt
69 lookAhead $ ensure len
70 if len == 0 then return KeepAlive
71 else do
72 mid <- getWord8
73 case mid of
74 0 -> return Choke
75 1 -> return Unchoke
76 2 -> return Interested
77 3 -> return NotInterested
78 4 -> Have <$> getInt
79 5 -> Bitfield <$> getBytes (pred len)
80 6 -> Request <$> get
81 7 -> Piece <$> getBlock (len - 9)
82 8 -> Cancel <$> get
83 9 -> (Port . fromIntegral) <$> getWord16be
84 _ -> fail $ "unknown message ID: " ++ show mid
85
86 where
87 getBlock :: Int -> Get Block
88 getBlock len = Block <$> getInt <*> getInt <*> getBytes len
89 {-# INLINE getBlock #-}
90
91 put KeepAlive = putInt 0
92 put Choke = putInt 1 >> putWord8 0
93 put Unchoke = putInt 1 >> putWord8 1
94 put Interested = putInt 1 >> putWord8 2
95 put NotInterested = putInt 1 >> putWord8 3
96 put (Have i) = putInt 5 >> putWord8 4 >> putInt i
97 put (Bitfield b) = putInt l >> putWord8 5 >> putByteString b
98 where l = succ (B.length b)
99 {-# INLINE l #-}
100 put (Request blk) = putInt 13 >> putWord8 6 >> put blk
101 put (Piece blk) = putInt l >> putWord8 7 >> putBlock
102 where l = 9 + B.length (blkData blk)
103 {-# INLINE l #-}
104 putBlock = do putInt (blkPiece blk)
105 putInt (blkOffset blk)
106 putByteString (blkData blk)
107 {-# INLINE putBlock #-}
108
109 put (Cancel blk) = putInt 13 >> putWord8 8 >> put blk
110 put (Port p ) = putInt 3 >> putWord8 9 >> putWord16be (fromIntegral p) \ No newline at end of file
diff --git a/src/Network/Torrent/PeerWire/Block.hs b/src/Network/Torrent/PeerWire/Block.hs
new file mode 100644
index 00000000..b2fe8c31
--- /dev/null
+++ b/src/Network/Torrent/PeerWire/Block.hs
@@ -0,0 +1,17 @@
1module Network.Torrent.PeerWire.Block
2 ( BlockIx(..), Block(..)
3 ) where
4
5import Data.ByteString (ByteString)
6
7data BlockIx = BlockIx {
8 ixPiece :: {-# UNPACK #-} !Int -- ^ Zero-based piece index.
9 , ixOffset :: {-# UNPACK #-} !Int -- ^ Zero-based byte offset within the piece.
10 , ixLength :: {-# UNPACK #-} !Int -- ^ Block size starting from offset.
11 } deriving (Show, Eq)
12
13data Block = Block {
14 blkPiece :: Int -- ^ Zero-based piece index.
15 , blkOffset :: Int -- ^ Zero-based byte offset within the piece.
16 , blkData :: ByteString -- ^ Payload.
17 } deriving (Show, Eq)
diff --git a/src/Network/Torrent/PeerWire/Message.hs b/src/Network/Torrent/PeerWire/Message.hs
new file mode 100644
index 00000000..e04038ff
--- /dev/null
+++ b/src/Network/Torrent/PeerWire/Message.hs
@@ -0,0 +1,91 @@
1module Network.Torrent.PeerWire.Message
2 ( Message(..)
3 ) where
4
5import Control.Applicative
6import Data.ByteString (ByteString)
7import qualified Data.ByteString as B
8
9import Data.Serialize
10
11import Network.Torrent.PeerWire.Block
12
13
14-- TODO comment message constructors
15data Message = KeepAlive
16 | Choke
17 | Unchoke
18 | Interested
19 | NotInterested
20 | Have Int
21 | Bitfield ByteString
22 | Request BlockIx
23 | Piece Block
24 | Cancel BlockIx
25 | Port Int
26 deriving (Show, Eq)
27
28instance Serialize BlockIx where
29 {-# SPECIALIZE instance Serialize BlockIx #-}
30 get = BlockIx <$> getInt <*> getInt <*> getInt
31 {-# INLINE get #-}
32
33 put ix = do putInt (ixPiece ix)
34 putInt (ixOffset ix)
35 putInt (ixLength ix)
36 {-# INLINE put #-}
37
38instance Serialize Message where
39 get = do
40 len <- getInt
41 lookAhead $ ensure len
42 if len == 0 then return KeepAlive -- FIX check if BS is empty instead of reading len
43 else do
44 mid <- getWord8
45 case mid of
46 0 -> return Choke
47 1 -> return Unchoke
48 2 -> return Interested
49 3 -> return NotInterested
50 4 -> Have <$> getInt
51 5 -> Bitfield <$> getBytes (pred len)
52 6 -> Request <$> get
53 7 -> Piece <$> getBlock (len - 9)
54 8 -> Cancel <$> get
55 9 -> (Port . fromIntegral) <$> getWord16be
56 _ -> fail $ "unknown message ID: " ++ show mid
57
58 where
59 getBlock :: Int -> Get Block
60 getBlock len = Block <$> getInt <*> getInt <*> getBytes len
61 {-# INLINE getBlock #-}
62
63 put KeepAlive = putInt 0
64 put Choke = putInt 1 >> putWord8 0
65 put Unchoke = putInt 1 >> putWord8 1
66 put Interested = putInt 1 >> putWord8 2
67 put NotInterested = putInt 1 >> putWord8 3
68 put (Have i) = putInt 5 >> putWord8 4 >> putInt i
69 put (Bitfield b) = putInt l >> putWord8 5 >> putByteString b
70 where l = succ (B.length b)
71 {-# INLINE l #-}
72 put (Request blk) = putInt 13 >> putWord8 6 >> put blk
73 put (Piece blk) = putInt l >> putWord8 7 >> putBlock
74 where l = 9 + B.length (blkData blk)
75 {-# INLINE l #-}
76 putBlock = do putInt (blkPiece blk)
77 putInt (blkOffset blk)
78 putByteString (blkData blk)
79 {-# INLINE putBlock #-}
80
81 put (Cancel blk) = putInt 13 >> putWord8 8 >> put blk
82 put (Port p ) = putInt 3 >> putWord8 9 >> putWord16be (fromIntegral p)
83
84
85getInt :: Get Int
86getInt = fromIntegral <$> getWord32be
87{-# INLINE getInt #-}
88
89putInt :: Putter Int
90putInt = putWord32be . fromIntegral
91{-# INLINE putInt #-}