3 files changed, 21 insertions, 2 deletions

diff --git a/src/Data/Torrent/Layout.hs b/src/Data/Torrent/Layout.hs
index 453c0d4f..a4b55b3d 100644
--- a/src/Data/Torrent/Layout.hs
+++ b/src/Data/Torrent/Layout.hs
@@ -56,6 +56,7 @@ module Data.Torrent.Layout
        , fileOffset
 
          -- * Internal
+       , sizeInBase
        , getLayoutInfo
        , putLayoutInfo
        ) where
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index bf44d7bf..8ad4d3e5 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -45,7 +45,7 @@ import Data.Typeable
 import Data.Torrent.Bitfield
 import Data.Torrent.Layout
 import Data.Torrent.Piece
-import System.Torrent.FileMap
+import System.Torrent.FileMap as FM
 
 
 data StorageFailure
@@ -82,7 +82,9 @@ withStorage :: Mode -> PieceSize -> FileLayout FileSize
 withStorage m s l = bracket (open m s l) close
 
 isValidIx :: PieceIx -> Storage -> Bool
-isValidIx i s = 0 <= i && i < undefined s
+isValidIx i Storage {..} = 0 <= i && i < pcount
+  where
+    pcount = FM.size fileMap `sizeInBase` pieceLen
 
 writePiece :: Piece BL.ByteString -> Storage -> IO ()
 writePiece p @ Piece {..} s @ Storage {..}
diff --git a/tests/System/Torrent/StorageSpec.hs b/tests/System/Torrent/StorageSpec.hs
index 8d9dfd8f..8267b7a5 100644
--- a/tests/System/Torrent/StorageSpec.hs
+++ b/tests/System/Torrent/StorageSpec.hs
@@ -1,5 +1,6 @@
 module System.Torrent.StorageSpec (spec) where
 import Control.Exception
+import Data.ByteString.Lazy as BL
 import System.FilePath
 import System.Directory
 import System.IO.Unsafe
@@ -32,11 +33,26 @@ spec = before createLayout $ do
       writePiece (Piece 0 "") s `shouldThrow` (== StorageIsRO)
       close s
 
+    it "should fail if piece size do not match" $ do
+      withStorage ReadWrite 1 layout $ \ s ->
+        writePiece (Piece 0 "") s `shouldThrow` (== InvalidSize 0)
+
     it "should fail on negative index" $ do
       withStorage ReadWrite 0 layout $ \ s ->
         writePiece (Piece (-1) "") s `shouldThrow` (== InvalidIndex (-1))
 
+    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)
+
   describe "readPiece" $ do
     it "should fail on negative index" $
       withStorage ReadOnly 0 layout $ \ s ->
         readPiece (-1) s `shouldThrow` (== InvalidIndex (-1))
+
+    it "should fail on out of upper bound index" $ do
+      withStorage ReadOnly 100 layout $ \ s -> do
+        _ <- readPiece 1 s
+        readPiece 2 s `shouldThrow` (== InvalidIndex 2)