Add storage streaming

3 files changed, 51 insertions, 0 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 2f30d848..67f1ef18 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -212,6 +212,7 @@ test-suite spec
 
                      , mtl
                      , resourcet
+                     , conduit
 
                      , hspec >= 1.8.1.1
                      , QuickCheck
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index 1c84bf69..3bb229a4 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -36,11 +36,17 @@ module System.Torrent.Storage
        , readPiece
        , hintRead
        , unsafeReadPiece
+
+         -- * Streaming
+       , sourceStorage
+       , sinkStorage
        ) where
 
 import Control.Applicative
 import Control.Exception
+import Control.Monad.Trans
 import Data.ByteString.Lazy as BL
+import Data.Conduit
 import Data.Typeable
 
 import Data.Torrent.Bitfield as BF
@@ -120,6 +126,24 @@ unsafeReadPiece pix s @ Storage {..}
     offset = fromIntegral pix * fromIntegral pieceLen
     sz     = fromIntegral pieceLen
 
+-- | Stream storage pieces from first to the last.
+sourceStorage :: Storage -> Source IO (Piece BL.ByteString)
+sourceStorage s = go 0
+  where
+    go pix
+        | pix < totalPieces s = do
+          piece <- liftIO $ readPiece pix s
+          liftIO $ hintRead (succ pix) s
+          yield piece
+          go (succ pix)
+        |       otherwise     = return ()
+
+-- | Write stream of pieces to the storage. Fail if storage is 'ReadOnly'.
+sinkStorage :: Storage -> Sink (Piece BL.ByteString) IO ()
+sinkStorage s = do
+  awaitForever $ \ piece ->
+    liftIO $ writePiece piece s
+
 -- | TODO examples of use
 genPieceInfo :: Storage -> IO PieceInfo
 genPieceInfo = undefined
diff --git a/tests/System/Torrent/StorageSpec.hs b/tests/System/Torrent/StorageSpec.hs
index 8267b7a5..d2185961 100644
--- a/tests/System/Torrent/StorageSpec.hs
+++ b/tests/System/Torrent/StorageSpec.hs
@@ -1,6 +1,8 @@
 module System.Torrent.StorageSpec (spec) where
 import Control.Exception
 import Data.ByteString.Lazy as BL
+import Data.Conduit as C
+import Data.Conduit.List as C
 import System.FilePath
 import System.Directory
 import System.IO.Unsafe
@@ -25,6 +27,12 @@ createLayout :: IO ()
 createLayout =
   bracket (open ReadWriteEx 0 layout) close (const (return ()))
 
+psize :: PieceSize
+psize = 16
+
+pcount :: PieceCount
+pcount = 11
+
 spec :: Spec
 spec = before createLayout $ do
   describe "writePiece" $ do
@@ -56,3 +64,21 @@ spec = before createLayout $ do
       withStorage ReadOnly 100 layout $ \ s -> do
         _ <- readPiece 1 s
         readPiece 2 s `shouldThrow` (== InvalidIndex 2)
+
+  describe "sourceStorage" $ do
+    it "should source all chunks" $ do
+      withStorage ReadOnly psize layout $ \ s -> do
+        n <- sourceStorage s $$ C.fold (\ n _ -> succ n) 0
+        n `shouldBe` pcount
+
+  -- this test should fail if 'sourceStorage' test fail
+  describe "sinkStorage" $ do
+    it "should write all chunks" $ do
+      let byteVal       = 0
+      let bzeroPiece  p = p { pieceData = BL.replicate (BL.length (pieceData p)) byteVal }
+      let isZeroPiece p = (== byteVal) `BL.all` pieceData p
+
+      withStorage ReadWrite psize layout $ \ s -> do
+        sourceStorage s $= C.map bzeroPiece $$ sinkStorage s
+        b <- sourceStorage s $$ C.fold (\ b p -> b && isZeroPiece p) True
+        b `shouldBe` True