diff options
-rw-r--r-- | bittorrent.cabal | 13 | ||||
-rw-r--r-- | src/Data/Bitfield.hs | 39 | ||||
-rw-r--r-- | src/Network/BitTorrent/PeerWire/Message.hs | 8 | ||||
-rw-r--r-- | src/Network/BitTorrent/Tracker.hs | 2 | ||||
-rw-r--r-- | tests/Main.hs | 12 |
5 files changed, 50 insertions, 24 deletions
diff --git a/bittorrent.cabal b/bittorrent.cabal index da69e48d..1a28c880 100644 --- a/bittorrent.cabal +++ b/bittorrent.cabal | |||
@@ -59,7 +59,8 @@ library | |||
59 | 59 | ||
60 | -- Data packages | 60 | -- Data packages |
61 | , array >= 0.4 | 61 | , array >= 0.4 |
62 | , bytestring >= 0.10.2 | 62 | , bytestring |
63 | -- >= 0.10.2 | ||
63 | , containers >= 0.4 | 64 | , containers >= 0.4 |
64 | , intset >= 0.1 | 65 | , intset >= 0.1 |
65 | , text >= 0.11.0 | 66 | , text >= 0.11.0 |
@@ -94,11 +95,12 @@ library | |||
94 | 95 | ||
95 | test-suite info-hash | 96 | test-suite info-hash |
96 | type: exitcode-stdio-1.0 | 97 | type: exitcode-stdio-1.0 |
97 | main-is: info-hash.hs | 98 | main-is: InfoHash.hs |
98 | hs-source-dirs: tests | 99 | hs-source-dirs: tests |
99 | 100 | ||
100 | build-depends: base == 4.* | 101 | build-depends: base == 4.* |
101 | , bytestring >= 0.10.2.0 | 102 | , bytestring |
103 | -- >= 0.10.2.0 | ||
102 | , containers >= 0.4.2.1 | 104 | , containers >= 0.4.2.1 |
103 | , bencoding >= 0.1.0.0 | 105 | , bencoding >= 0.1.0.0 |
104 | , bittorrent | 106 | , bittorrent |
@@ -116,7 +118,8 @@ test-suite properties | |||
116 | main-is: Main.hs | 118 | main-is: Main.hs |
117 | hs-source-dirs: tests | 119 | hs-source-dirs: tests |
118 | build-depends: base == 4.* | 120 | build-depends: base == 4.* |
119 | , bytestring >= 0.10.2 | 121 | , bytestring |
122 | -- >= 0.10.2 | ||
120 | , cereal >= 0.3.5.2 | 123 | , cereal >= 0.3.5.2 |
121 | , network >= 2.4.0.13 | 124 | , network >= 2.4.0.13 |
122 | , text | 125 | , text |
@@ -141,7 +144,7 @@ benchmark benchmarks | |||
141 | hs-source-dirs: bench | 144 | hs-source-dirs: bench |
142 | 145 | ||
143 | build-depends: base == 4.* | 146 | build-depends: base == 4.* |
144 | , bytestring >= 0.10.2.0 | 147 | , bytestring |
145 | , cereal | 148 | , cereal |
146 | , network | 149 | , network |
147 | 150 | ||
diff --git a/src/Data/Bitfield.hs b/src/Data/Bitfield.hs index ee0570e7..56365bf7 100644 --- a/src/Data/Bitfield.hs +++ b/src/Data/Bitfield.hs | |||
@@ -17,6 +17,7 @@ module Data.Bitfield | |||
17 | 17 | ||
18 | -- * Construction | 18 | -- * Construction |
19 | , haveAll, haveNone, have | 19 | , haveAll, haveNone, have |
20 | , adjustSize | ||
20 | 21 | ||
21 | -- * Query | 22 | -- * Query |
22 | , Data.Bitfield.null | 23 | , Data.Bitfield.null |
@@ -31,8 +32,7 @@ module Data.Bitfield | |||
31 | , difference | 32 | , difference |
32 | 33 | ||
33 | -- * Serialization | 34 | -- * Serialization |
34 | , getBitfield, putBitfield | 35 | , fromBitmap, toBitmap |
35 | , bitfieldByteCount | ||
36 | 36 | ||
37 | #if defined (TESTING) | 37 | #if defined (TESTING) |
38 | -- * Debug | 38 | -- * Debug |
@@ -42,15 +42,18 @@ module Data.Bitfield | |||
42 | 42 | ||
43 | import Control.Monad | 43 | import Control.Monad |
44 | import Control.Monad.ST | 44 | import Control.Monad.ST |
45 | import Data.ByteString (ByteString) | ||
46 | import qualified Data.ByteString as B | ||
47 | import qualified Data.ByteString.Lazy as Lazy | ||
45 | import Data.Vector.Unboxed (Vector) | 48 | import Data.Vector.Unboxed (Vector) |
46 | import qualified Data.Vector.Unboxed as V | 49 | import qualified Data.Vector.Unboxed as V |
47 | import qualified Data.Vector.Unboxed.Mutable as VM | 50 | import qualified Data.Vector.Unboxed.Mutable as VM |
48 | import Data.IntervalSet (IntSet) | 51 | import Data.IntervalSet (IntSet) |
49 | import qualified Data.IntervalSet as S | 52 | import qualified Data.IntervalSet as S |
53 | import qualified Data.IntervalSet.ByteString as S | ||
50 | import Data.List (foldl') | 54 | import Data.List (foldl') |
51 | import Data.Monoid | 55 | import Data.Monoid |
52 | import Data.Ratio | 56 | import Data.Ratio |
53 | import Data.Serialize | ||
54 | import Network.BitTorrent.PeerWire.Block | 57 | import Network.BitTorrent.PeerWire.Block |
55 | 58 | ||
56 | 59 | ||
@@ -97,6 +100,11 @@ have ix Bitfield {..} | |||
97 | | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet) | 100 | | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet) |
98 | | otherwise = Bitfield bfSize bfSet | 101 | | otherwise = Bitfield bfSize bfSet |
99 | 102 | ||
103 | -- | Assign new size to bitfield. FIXME Normally, size should be only | ||
104 | -- decreased, otherwise exception raised. | ||
105 | adjustSize :: PieceCount -> Bitfield -> Bitfield | ||
106 | adjustSize s Bitfield {..} = Bitfield s bfSet | ||
107 | |||
100 | {----------------------------------------------------------------------- | 108 | {----------------------------------------------------------------------- |
101 | Query | 109 | Query |
102 | -----------------------------------------------------------------------} | 110 | -----------------------------------------------------------------------} |
@@ -200,17 +208,22 @@ unions = foldl' union (haveNone 0) | |||
200 | Serialization | 208 | Serialization |
201 | -----------------------------------------------------------------------} | 209 | -----------------------------------------------------------------------} |
202 | 210 | ||
203 | -- | | 211 | -- | Unpack 'Bitfield' from tightly packed bit array. Note resulting |
204 | getBitfield :: Int -> Get Bitfield | 212 | -- size might be more than real bitfield size, use 'adjustSize'. |
205 | getBitfield = error "getBitfield" | 213 | fromBitmap :: ByteString -> Bitfield |
206 | 214 | fromBitmap bs = Bitfield { | |
207 | -- | | 215 | bfSize = B.length bs * 8 |
208 | putBitfield :: Bitfield -> Put | 216 | , bfSet = S.fromByteString bs |
209 | putBitfield = error "putBitfield" | 217 | } |
218 | {-# INLINE fromBitmap #-} | ||
210 | 219 | ||
211 | -- | | 220 | -- | Pack a 'Bitfield' to tightly packed bit array. |
212 | bitfieldByteCount :: Bitfield -> Int | 221 | toBitmap :: Bitfield -> Lazy.ByteString |
213 | bitfieldByteCount = error "bitfieldByteCount" | 222 | toBitmap Bitfield {..} = Lazy.fromChunks [intsetBM, alignment] |
223 | where | ||
224 | byteSize = bfSize `div` 8 + if bfSize `mod` 8 == 0 then 0 else 1 | ||
225 | alignment = B.replicate (byteSize - B.length intsetBM) 0 | ||
226 | intsetBM = S.toByteString bfSet | ||
214 | 227 | ||
215 | {----------------------------------------------------------------------- | 228 | {----------------------------------------------------------------------- |
216 | Debug | 229 | Debug |
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs index 6515fdf2..3895ed5f 100644 --- a/src/Network/BitTorrent/PeerWire/Message.hs +++ b/src/Network/BitTorrent/PeerWire/Message.hs | |||
@@ -7,6 +7,7 @@ module Network.BitTorrent.PeerWire.Message | |||
7 | 7 | ||
8 | import Control.Applicative | 8 | import Control.Applicative |
9 | import qualified Data.ByteString as B | 9 | import qualified Data.ByteString as B |
10 | import qualified Data.ByteString.Lazy as Lazy | ||
10 | import Data.Serialize | 11 | import Data.Serialize |
11 | import Text.PrettyPrint | 12 | import Text.PrettyPrint |
12 | import Network | 13 | import Network |
@@ -91,7 +92,7 @@ instance Serialize Message where | |||
91 | 0x02 -> return Interested | 92 | 0x02 -> return Interested |
92 | 0x03 -> return NotInterested | 93 | 0x03 -> return NotInterested |
93 | 0x04 -> Have <$> getInt | 94 | 0x04 -> Have <$> getInt |
94 | 0x05 -> Bitfield <$> getBitfield (pred len) | 95 | 0x05 -> (Bitfield . fromBitmap) <$> getByteString (pred len) |
95 | 0x06 -> Request <$> get | 96 | 0x06 -> Request <$> get |
96 | 0x07 -> Piece <$> getBlock (len - 9) | 97 | 0x07 -> Piece <$> getBlock (len - 9) |
97 | 0x08 -> Cancel <$> get | 98 | 0x08 -> Cancel <$> get |
@@ -118,8 +119,9 @@ instance Serialize Message where | |||
118 | put Interested = putInt 1 >> putWord8 0x02 | 119 | put Interested = putInt 1 >> putWord8 0x02 |
119 | put NotInterested = putInt 1 >> putWord8 0x03 | 120 | put NotInterested = putInt 1 >> putWord8 0x03 |
120 | put (Have i) = putInt 5 >> putWord8 0x04 >> putInt i | 121 | put (Have i) = putInt 5 >> putWord8 0x04 >> putInt i |
121 | put (Bitfield b) = putInt l >> putWord8 0x05 >> putBitfield b | 122 | put (Bitfield bf) = putInt l >> putWord8 0x05 >> putLazyByteString b |
122 | where l = succ (bitfieldByteCount b) | 123 | where b = toBitmap bf |
124 | l = succ (fromIntegral (Lazy.length b)) | ||
123 | {-# INLINE l #-} | 125 | {-# INLINE l #-} |
124 | put (Request blk) = putInt 13 >> putWord8 0x06 >> put blk | 126 | put (Request blk) = putInt 13 >> putWord8 0x06 >> put blk |
125 | put (Piece blk) = putInt l >> putWord8 0x07 >> putBlock | 127 | put (Piece blk) = putInt l >> putWord8 0x07 >> putBlock |
diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs index 2e599002..aaa08f3c 100644 --- a/src/Network/BitTorrent/Tracker.hs +++ b/src/Network/BitTorrent/Tracker.hs | |||
@@ -145,7 +145,7 @@ data TSession = TSession { | |||
145 | newSession :: Progress -> Int -> [PeerAddr] -> IO TSession | 145 | newSession :: Progress -> Int -> [PeerAddr] -> IO TSession |
146 | newSession pr i ps = TSession <$> newTVarIO pr | 146 | newSession pr i ps = TSession <$> newTVarIO pr |
147 | <*> newIORef i | 147 | <*> newIORef i |
148 | <*> newTVarIO psx | 148 | <*> newTVarIO ps |
149 | 149 | ||
150 | getPeerList :: TSession -> IO [PeerAddr] | 150 | getPeerList :: TSession -> IO [PeerAddr] |
151 | getPeerList = readTVarIO . sePeers | 151 | getPeerList = readTVarIO . sePeers |
diff --git a/tests/Main.hs b/tests/Main.hs index 0e18a06b..84870d66 100644 --- a/tests/Main.hs +++ b/tests/Main.hs | |||
@@ -185,6 +185,15 @@ instance Arbitrary Message where | |||
185 | , Cancel <$> arbitrary | 185 | , Cancel <$> arbitrary |
186 | , Port <$> arbitrary | 186 | , Port <$> arbitrary |
187 | ] | 187 | ] |
188 | -- todo add all messages | ||
189 | |||
190 | prop_messageEncoding :: Message -> Bool | ||
191 | prop_messageEncoding msg @ (Bitfield bf) | ||
192 | = case S.decode (S.encode msg) of | ||
193 | Right (Bitfield bf') -> bf == adjustSize (totalCount bf) bf' | ||
194 | Left _ -> False | ||
195 | prop_messageEncoding msg | ||
196 | = S.decode (S.encode msg) == Right msg | ||
188 | 197 | ||
189 | {----------------------------------------------------------------------- | 198 | {----------------------------------------------------------------------- |
190 | Main | 199 | Main |
@@ -209,7 +218,6 @@ main = defaultMain $ | |||
209 | -- handshake module | 218 | -- handshake module |
210 | , testProperty "handshake encoding" $ | 219 | , testProperty "handshake encoding" $ |
211 | prop_cerealEncoding (T :: T Handshake) | 220 | prop_cerealEncoding (T :: T Handshake) |
212 | , testProperty "message encoding" $ | 221 | , testProperty "message encoding" prop_messageEncoding |
213 | prop_cerealEncoding (T :: T Message) | ||
214 | 222 | ||
215 | ] ++ test_scrape_url | 223 | ] ++ test_scrape_url |