diff options
-rw-r--r-- | network-bittorrent.cabal | 3 | ||||
-rw-r--r-- | src/Network/Torrent/PeerWire.hs | 103 | ||||
-rw-r--r-- | src/Network/Torrent/PeerWire/Block.hs | 17 | ||||
-rw-r--r-- | src/Network/Torrent/PeerWire/Message.hs | 91 |
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 #-} |
9 | module Network.Torrent.PeerWire | 9 | module 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 | ||
15 | import Network.Torrent.PeerWire.Block | ||
16 | import Network.Torrent.PeerWire.Message | ||
14 | import Network.Torrent.PeerWire.Handshake | 17 | import Network.Torrent.PeerWire.Handshake |
15 | |||
16 | import Control.Applicative | ||
17 | import Data.ByteString (ByteString) | ||
18 | import qualified Data.ByteString as B | ||
19 | |||
20 | import Data.Serialize | ||
21 | |||
22 | data 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 | |||
28 | data 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 | ||
35 | data 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 | |||
48 | getInt :: Get Int | ||
49 | getInt = fromIntegral <$> getWord32be | ||
50 | {-# INLINE getInt #-} | ||
51 | |||
52 | putInt :: Putter Int | ||
53 | putInt = putWord32be . fromIntegral | ||
54 | {-# INLINE putInt #-} | ||
55 | |||
56 | instance 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 | |||
66 | instance 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 @@ | |||
1 | module Network.Torrent.PeerWire.Block | ||
2 | ( BlockIx(..), Block(..) | ||
3 | ) where | ||
4 | |||
5 | import Data.ByteString (ByteString) | ||
6 | |||
7 | data 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 | |||
13 | data 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 @@ | |||
1 | module Network.Torrent.PeerWire.Message | ||
2 | ( Message(..) | ||
3 | ) where | ||
4 | |||
5 | import Control.Applicative | ||
6 | import Data.ByteString (ByteString) | ||
7 | import qualified Data.ByteString as B | ||
8 | |||
9 | import Data.Serialize | ||
10 | |||
11 | import Network.Torrent.PeerWire.Block | ||
12 | |||
13 | |||
14 | -- TODO comment message constructors | ||
15 | data 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 | |||
28 | instance 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 | |||
38 | instance 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 | |||
85 | getInt :: Get Int | ||
86 | getInt = fromIntegral <$> getWord32be | ||
87 | {-# INLINE getInt #-} | ||
88 | |||
89 | putInt :: Putter Int | ||
90 | putInt = putWord32be . fromIntegral | ||
91 | {-# INLINE putInt #-} | ||