diff options
-rw-r--r-- | examples/Main.hs | 2 | ||||
-rw-r--r-- | src/Network/BitTorrent/Exchange/Protocol.hs | 26 | ||||
-rw-r--r-- | src/System/Torrent/Storage.hs | 9 | ||||
-rw-r--r-- | tests/Main.hs | 3 |
4 files changed, 23 insertions, 17 deletions
diff --git a/examples/Main.hs b/examples/Main.hs index 18cbefe3..edb37975 100644 --- a/examples/Main.hs +++ b/examples/Main.hs | |||
@@ -13,7 +13,7 @@ main = do | |||
13 | 13 | ||
14 | print (contentLayout "./" (tInfo torrent)) | 14 | print (contentLayout "./" (tInfo torrent)) |
15 | 15 | ||
16 | client <- newClient 10 [] | 16 | client <- newClient 2 [] |
17 | swarm <- newLeecher client torrent | 17 | swarm <- newLeecher client torrent |
18 | 18 | ||
19 | storage <- swarm `bindTo` "/tmp/" | 19 | storage <- swarm `bindTo` "/tmp/" |
diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs index a285f8d2..feccf760 100644 --- a/src/Network/BitTorrent/Exchange/Protocol.hs +++ b/src/Network/BitTorrent/Exchange/Protocol.hs | |||
@@ -260,7 +260,7 @@ data Block = Block { | |||
260 | , blkOffset :: {-# UNPACK #-} !Int | 260 | , blkOffset :: {-# UNPACK #-} !Int |
261 | 261 | ||
262 | -- | Payload. | 262 | -- | Payload. |
263 | , blkData :: !ByteString -- TODO make lazy bytestring | 263 | , blkData :: !Lazy.ByteString -- TODO make lazy bytestring |
264 | } deriving (Show, Eq) | 264 | } deriving (Show, Eq) |
265 | 265 | ||
266 | -- | Format block in human readable form. Payload is ommitted. | 266 | -- | Format block in human readable form. Payload is ommitted. |
@@ -268,7 +268,8 @@ ppBlock :: Block -> Doc | |||
268 | ppBlock = ppBlockIx . blockIx | 268 | ppBlock = ppBlockIx . blockIx |
269 | 269 | ||
270 | blockSize :: Block -> Int | 270 | blockSize :: Block -> Int |
271 | blockSize blk = B.length (blkData blk) | 271 | blockSize blk = fromIntegral (Lazy.length (blkData blk)) |
272 | {-# INLINE blockSize #-} | ||
272 | 273 | ||
273 | -- | Widely used semi-official block size. | 274 | -- | Widely used semi-official block size. |
274 | defaultBlockSize :: Int | 275 | defaultBlockSize :: Int |
@@ -277,7 +278,9 @@ defaultBlockSize = 16 * 1024 | |||
277 | 278 | ||
278 | isPiece :: Int -> Block -> Bool | 279 | isPiece :: Int -> Block -> Bool |
279 | isPiece pieceSize (Block i offset bs) = | 280 | isPiece pieceSize (Block i offset bs) = |
280 | offset == 0 && B.length bs == pieceSize && i >= 0 | 281 | offset == 0 |
282 | && fromIntegral (Lazy.length bs) == pieceSize | ||
283 | && i >= 0 | ||
281 | {-# INLINE isPiece #-} | 284 | {-# INLINE isPiece #-} |
282 | 285 | ||
283 | pieceIx :: Int -> Int -> BlockIx | 286 | pieceIx :: Int -> Int -> BlockIx |
@@ -285,14 +288,14 @@ pieceIx i = BlockIx i 0 | |||
285 | {-# INLINE pieceIx #-} | 288 | {-# INLINE pieceIx #-} |
286 | 289 | ||
287 | blockIx :: Block -> BlockIx | 290 | blockIx :: Block -> BlockIx |
288 | blockIx = BlockIx <$> blkPiece <*> blkOffset <*> B.length . blkData | 291 | blockIx = BlockIx <$> blkPiece <*> blkOffset <*> blockSize |
289 | 292 | ||
290 | blockRange :: (Num a, Integral a) => Int -> Block -> (a, a) | 293 | blockRange :: (Num a, Integral a) => Int -> Block -> (a, a) |
291 | blockRange pieceSize blk = (offset, offset + len) | 294 | blockRange pieceSize blk = (offset, offset + len) |
292 | where | 295 | where |
293 | offset = fromIntegral pieceSize * fromIntegral (blkPiece blk) | 296 | offset = fromIntegral pieceSize * fromIntegral (blkPiece blk) |
294 | + fromIntegral (blkOffset blk) | 297 | + fromIntegral (blkOffset blk) |
295 | len = fromIntegral (B.length (blkData blk)) | 298 | len = fromIntegral (Lazy.length (blkData blk)) |
296 | {-# INLINE blockRange #-} | 299 | {-# INLINE blockRange #-} |
297 | {-# SPECIALIZE blockRange :: Int -> Block -> (Int64, Int64) #-} | 300 | {-# SPECIALIZE blockRange :: Int -> Block -> (Int64, Int64) #-} |
298 | 301 | ||
@@ -402,7 +405,7 @@ instance Serialize Message where | |||
402 | 405 | ||
403 | where | 406 | where |
404 | getBlock :: Int -> S.Get Block | 407 | getBlock :: Int -> S.Get Block |
405 | getBlock len = Block <$> getInt <*> getInt <*> S.getBytes len | 408 | getBlock len = Block <$> getInt <*> getInt <*> S.getLazyByteString (fromIntegral len) |
406 | {-# INLINE getBlock #-} | 409 | {-# INLINE getBlock #-} |
407 | 410 | ||
408 | 411 | ||
@@ -418,11 +421,11 @@ instance Serialize Message where | |||
418 | {-# INLINE l #-} | 421 | {-# INLINE l #-} |
419 | put (Request blk) = putInt 13 >> S.putWord8 0x06 >> S.put blk | 422 | put (Request blk) = putInt 13 >> S.putWord8 0x06 >> S.put blk |
420 | put (Piece blk) = putInt l >> S.putWord8 0x07 >> putBlock | 423 | put (Piece blk) = putInt l >> S.putWord8 0x07 >> putBlock |
421 | where l = 9 + B.length (blkData blk) | 424 | where l = 9 + fromIntegral (Lazy.length (blkData blk)) |
422 | {-# INLINE l #-} | 425 | {-# INLINE l #-} |
423 | putBlock = do putInt (blkPiece blk) | 426 | putBlock = do putInt (blkPiece blk) |
424 | putInt (blkOffset blk) | 427 | putInt (blkOffset blk) |
425 | S.putByteString (blkData blk) | 428 | S.putLazyByteString (blkData blk) |
426 | {-# INLINE putBlock #-} | 429 | {-# INLINE putBlock #-} |
427 | 430 | ||
428 | put (Cancel blk) = putInt 13 >> S.putWord8 0x08 >> S.put blk | 431 | put (Cancel blk) = putInt 13 >> S.putWord8 0x08 >> S.put blk |
@@ -463,7 +466,8 @@ instance Binary Message where | |||
463 | 466 | ||
464 | where | 467 | where |
465 | getBlock :: Int -> B.Get Block | 468 | getBlock :: Int -> B.Get Block |
466 | getBlock len = Block <$> getIntB <*> getIntB <*> B.getByteString len | 469 | getBlock len = Block <$> getIntB <*> getIntB |
470 | <*> B.getLazyByteString (fromIntegral len) | ||
467 | {-# INLINE getBlock #-} | 471 | {-# INLINE getBlock #-} |
468 | 472 | ||
469 | put KeepAlive = putIntB 0 | 473 | put KeepAlive = putIntB 0 |
@@ -478,11 +482,11 @@ instance Binary Message where | |||
478 | {-# INLINE l #-} | 482 | {-# INLINE l #-} |
479 | put (Request blk) = putIntB 13 >> B.putWord8 0x06 >> B.put blk | 483 | put (Request blk) = putIntB 13 >> B.putWord8 0x06 >> B.put blk |
480 | put (Piece blk) = putIntB l >> B.putWord8 0x07 >> putBlock | 484 | put (Piece blk) = putIntB l >> B.putWord8 0x07 >> putBlock |
481 | where l = 9 + B.length (blkData blk) | 485 | where l = 9 + fromIntegral (Lazy.length (blkData blk)) |
482 | {-# INLINE l #-} | 486 | {-# INLINE l #-} |
483 | putBlock = do putIntB (blkPiece blk) | 487 | putBlock = do putIntB (blkPiece blk) |
484 | putIntB (blkOffset blk) | 488 | putIntB (blkOffset blk) |
485 | B.putByteString (blkData blk) | 489 | B.putLazyByteString (blkData blk) |
486 | {-# INLINE putBlock #-} | 490 | {-# INLINE putBlock #-} |
487 | 491 | ||
488 | put (Cancel blk) = putIntB 13 >> B.putWord8 0x08 >> B.put blk | 492 | put (Cancel blk) = putIntB 13 >> B.putWord8 0x08 >> B.put blk |
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs index 955c1746..363f94ef 100644 --- a/src/System/Torrent/Storage.hs +++ b/src/System/Torrent/Storage.hs | |||
@@ -143,9 +143,7 @@ putBlk blk @ Block {..} st @ Storage {..} | |||
143 | -- let blkIx = undefined | 143 | -- let blkIx = undefined |
144 | -- bm <- readTVarIO blocks | 144 | -- bm <- readTVarIO blocks |
145 | -- unless (member blkIx bm) $ do | 145 | -- unless (member blkIx bm) $ do |
146 | writeBytes (blkInterval (pieceLength session) blk) | 146 | writeBytes (blkInterval (pieceLength session) blk) blkData payload |
147 | (Lazy.fromChunks [blkData]) | ||
148 | payload | ||
149 | 147 | ||
150 | markBlock blk st | 148 | markBlock blk st |
151 | validatePiece blkPiece st | 149 | validatePiece blkPiece st |
@@ -166,7 +164,7 @@ getBlk ix @ BlockIx {..} st @ Storage {..} | |||
166 | = liftIO $ {-# SCC getBlk #-} do | 164 | = liftIO $ {-# SCC getBlk #-} do |
167 | -- TODO check if __piece__ is available | 165 | -- TODO check if __piece__ is available |
168 | bs <- readBytes (ixInterval (pieceLength session) ix) payload | 166 | bs <- readBytes (ixInterval (pieceLength session) ix) payload |
169 | return $ Block ixPiece ixOffset (Lazy.toStrict bs) | 167 | return $ Block ixPiece ixOffset bs |
170 | 168 | ||
171 | getPiece :: PieceIx -> Storage -> IO ByteString | 169 | getPiece :: PieceIx -> Storage -> IO ByteString |
172 | getPiece pix st @ Storage {..} = {-# SCC getPiece #-} do | 170 | getPiece pix st @ Storage {..} = {-# SCC getPiece #-} do |
@@ -220,4 +218,5 @@ ixInterval pieceSize BlockIx {..} = | |||
220 | 218 | ||
221 | blkInterval :: Int -> Block -> FixedInterval | 219 | blkInterval :: Int -> Block -> FixedInterval |
222 | blkInterval pieceSize Block {..} = | 220 | blkInterval pieceSize Block {..} = |
223 | Fixed.interval (blkPiece * pieceSize + blkOffset) (B.length blkData) \ No newline at end of file | 221 | Fixed.interval (blkPiece * pieceSize + blkOffset) |
222 | (fromIntegral (Lazy.length blkData)) \ No newline at end of file | ||
diff --git a/tests/Main.hs b/tests/Main.hs index b99f2469..45f92813 100644 --- a/tests/Main.hs +++ b/tests/Main.hs | |||
@@ -161,6 +161,9 @@ positive = fromIntegral <$> (arbitrary :: Gen Word32) | |||
161 | instance Arbitrary ByteString where | 161 | instance Arbitrary ByteString where |
162 | arbitrary = B.pack <$> arbitrary | 162 | arbitrary = B.pack <$> arbitrary |
163 | 163 | ||
164 | instance Arbitrary Lazy.ByteString where | ||
165 | arbitrary = Lazy.pack <$> arbitrary | ||
166 | |||
164 | instance Arbitrary BlockIx where | 167 | instance Arbitrary BlockIx where |
165 | arbitrary = BlockIx <$> positive <*> positive <*> positive | 168 | arbitrary = BlockIx <$> positive <*> positive <*> positive |
166 | 169 | ||