From f1d79ea2180e46545f442985a85e5bb1cce90436 Mon Sep 17 00:00:00 2001 From: Sam Truzjan Date: Sun, 5 Jan 2014 20:19:01 +0400 Subject: Check size of last piece properly (in writePiece) --- src/System/Torrent/Storage.hs | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) (limited to 'src/System/Torrent') diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs index 3bb229a4..8aa1aa99 100644 --- a/src/System/Torrent/Storage.hs +++ b/src/System/Torrent/Storage.hs @@ -96,11 +96,23 @@ isValidIx i s = 0 <= i && i < totalPieces s writePiece :: Piece BL.ByteString -> Storage -> IO () writePiece p @ Piece {..} s @ Storage {..} - | mode == ReadOnly = throwIO StorageIsRO - | pieceSize p /= pieceLen = throwIO (InvalidSize (pieceSize p)) - | not (isValidIx pieceIndex s) = throwIO (InvalidIndex pieceIndex) - | otherwise = writeBytes offset pieceData fileMap + | mode == ReadOnly = throwIO StorageIsRO + | isNotValidIx pieceIndex = throwIO (InvalidIndex pieceIndex) + | isNotValidSize pieceIndex (pieceSize p) + = throwIO (InvalidSize (pieceSize p)) + | otherwise = writeBytes offset pieceData fileMap where + isNotValidSize pix psize + | succ pix == pcount = psize /= lastPieceLen -- last piece may be shorter + | otherwise = psize /= pieceLen + where + lastPieceLen = fromIntegral (FM.size fileMap `rem` fromIntegral pieceLen) + {-# INLINE isNotValidSize #-} + + isNotValidIx i = i < 0 || i >= pcount + {-# INLINE isNotValidIx #-} + + pcount = totalPieces s offset = fromIntegral pieceIndex * fromIntegral pieceLen readPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString) -- cgit v1.2.3