2 files changed, 22 insertions, 6 deletions

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)
diff --git a/tests/System/Torrent/StorageSpec.hs b/tests/System/Torrent/StorageSpec.hs
index d2185961..2e94ccf1 100644
--- a/tests/System/Torrent/StorageSpec.hs
+++ b/tests/System/Torrent/StorageSpec.hs
@@ -52,8 +52,12 @@ spec = before createLayout $ do
     it "should fail on out of upper bound index" $ do
       withStorage ReadWrite 100 layout $ \ s -> do
         let bs = BL.replicate 100 0
-        writePiece (Piece 1 bs) s
-        writePiece (Piece 2 bs) s `shouldThrow` (== InvalidIndex 2)
+        writePiece (Piece 0 bs) s
+
+        let bs' = BL.replicate 75 0
+        writePiece (Piece 1 bs') s
+
+        writePiece (Piece 2 bs') s `shouldThrow` (== InvalidIndex 2)
 
   describe "readPiece" $ do
     it "should fail on negative index" $