Add check command to mktorrent utility

4 files changed, 84 insertions, 20 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index d721eacd..073329de 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -261,13 +261,13 @@ executable mktorrent
                      , pretty-class
 
                      , lens
-
+                     , parallel-io
                      , network
                      , bittorrent
 
                      , optparse-applicative
                      , hslogger
-  ghc-options:         -Wall -threaded
+  ghc-options:         -Wall -O2 -threaded
 
 -- Utility to fetch
 executable gettorrent
diff --git a/examples/MkTorrent.hs b/examples/MkTorrent.hs
index 7bf8b513..49856f9b 100644
--- a/examples/MkTorrent.hs
+++ b/examples/MkTorrent.hs
@@ -5,10 +5,12 @@
 module Main (main) where
 
 import Prelude as P
+import Control.Concurrent.ParallelIO
 import Control.Exception
 import Control.Lens hiding (argument)
 import Control.Monad
 import Data.List as L
+import Data.Maybe
 import Data.Monoid
 import Data.Text as T
 import qualified Data.Text.IO as T
@@ -24,8 +26,12 @@ import Text.PrettyPrint.Class
 
 import Paths_bittorrent (version)
 import Data.Torrent
+import Data.Torrent.Bitfield as BF
+import Data.Torrent.Piece
+import Data.Torrent.Layout
 import Data.Torrent.Magnet hiding (Magnet (Magnet))
 import Data.Torrent.Magnet (Magnet)
+import System.Torrent.Storage
 
 
 {-----------------------------------------------------------------------
@@ -138,27 +144,66 @@ amend (AmendOpts tpath) = do
   toFile tpath $ a t
 
 {-----------------------------------------------------------------------
---  Check command
+--  Check command -- validate content files using torrent file
 -----------------------------------------------------------------------}
 
-{-
-checkOpts :: Parser CheckOpts
-checkOpts = CheckOpts
-    <$> torrentFile
-    <*> argument Just
-        ( metavar "PATH"
-       <> value   "."
-       <> help    "Content directory or a single file" )
+data CheckOpts = CheckOpts
+  { checkTorrentPath :: FilePath -- ^ validation torrent file
+  , checkContentPath :: FilePath -- ^ root dir for content files
+  } deriving Show
 
 checkInfo :: ParserInfo CheckOpts
-checkInfo = info (helper <*> checkOpts) modifier
+checkInfo = info (helper <*> parser) modifier
   where
     modifier = progDesc "Validate integrity of torrent data"
+    parser   = CheckOpts
+      <$> torrentFile
+      <*> argument Just
+          ( metavar "PATH"
+         <> value   "."
+         <> help    "Content directory or a single file"
+          )
+
+validatePiece :: Storage -> PieceInfo -> PieceIx -> IO (Maybe PieceIx)
+validatePiece s pinfo pix = do
+  valid <- verifyPiece s pinfo pix
+  if valid
+    then do infoM "check" $ "valid piece "   ++ show pix
+            return (Just pix)
+    else do infoM "check" $ "invalid piece " ++ show pix
+            return Nothing
+
+validateStorage :: Storage -> PieceInfo -> IO Bitfield
+validateStorage s pinfo = do
+  infoM "check" "start storage validation"
+  let total = totalPieces s
+  pixs <- parallel $ L.map (validatePiece s pinfo) [0 .. total - 1]
+  infoM "check" "storage validation finished"
+  return $ fromList total $ catMaybes pixs
+
+-- TODO use local thread pool
+checkContent :: Storage -> PieceInfo -> IO ()
+checkContent s pinfo = do
+  invalids <- BF.complement <$> validateStorage s pinfo
+  if BF.null invalids
+    then noticeM "check" "all files are complete and valid"
+    else do
+      emergencyM "check" $ "there are some invalid pieces" ++ show invalids
+      exitFailure
+
+checkTorrent :: CheckOpts -> IO ()
+checkTorrent CheckOpts {..} = do
+  InfoDict {..} <- tInfoDict <$> fromFile checkTorrentPath
+  let layout = flatLayout checkContentPath idLayoutInfo
+  withStorage ReadOnly (piPieceLength idPieceInfo) layout $ \ s -> do
+    infoM "check" "files mapped"
+    checkContent s idPieceInfo
+    infoM "check" "unmapping files"
 
 {-----------------------------------------------------------------------
 --  Create command
 -----------------------------------------------------------------------}
-
+{-
 createFlags :: Parser CreateFlags
 createFlags = CreateFlags
     <$> optional (option
@@ -257,7 +302,7 @@ putTorrent opts @ ShowOpts {..} = do
 
 data Command
   = Amend  AmendOpts
---  | Check  CheckOpts
+  | Check  CheckOpts
 --  | Create CreateOpts
   | Magnet MagnetOpts
   | Show   ShowOpts
@@ -266,7 +311,7 @@ data Command
 commandOpts :: Parser Command
 commandOpts = subparser $ mconcat
     [ command "amend"  (Amend  <$> amendInfo)
---    , command "check"  (Check  <$> checkInfo)
+    , command "check"  (Check  <$> checkInfo)
 --    , command "create" (Create <$> createInfo)
     , command "magnet" (Magnet <$> magnetInfo)
     , command "show"   (Show   <$> showInfo)
@@ -349,7 +394,7 @@ parserInfo = info parser modifier
 
 run :: Command -> IO ()
 run (Amend  opts) = amend opts
---run (Check  opts) = checkTorrent opts
+run (Check  opts) = checkTorrent opts
 --run (Create opts) = createTorrent opts
 run (Magnet opts) = magnet opts
 run (Show   opts) = putTorrent opts
diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index 8cdae69f..b65f058b 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -57,6 +57,7 @@ module Data.Torrent.Bitfield
        , isSubsetOf
 
          -- ** Availability
+       , complement
        , Frequency
        , frequencies
        , rarest
@@ -194,8 +195,20 @@ findMax :: Bitfield -> PieceIx
 findMax = S.findMax . bfSet
 {-# INLINE findMax #-}
 
+-- | Check if all pieces from first bitfield present if the second bitfield
 isSubsetOf :: Bitfield -> Bitfield -> Bool
 isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b
+{-# INLINE isSubsetOf #-}
+
+-- | Resulting bitfield includes only missing pieces.
+complement :: Bitfield -> Bitfield
+complement Bitfield {..} = Bitfield
+  { bfSet  = uni `S.difference` bfSet
+  , bfSize = bfSize
+  }
+  where
+    Bitfield _ uni = haveAll bfSize
+{-# INLINE complement #-}
 
 {-----------------------------------------------------------------------
 -- Availability
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index a8b0bdc6..b5092b2e 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -28,6 +28,8 @@ module System.Torrent.Storage
        , withStorage
 
          -- * Query
+       , totalPieces
+       , verifyPiece
        , genPieceInfo
        , getBitfield
 
@@ -166,6 +168,10 @@ genPieceInfo s = do
   hashes <- sourceStorage s $= C.map hashPiece $$ C.sinkLbs
   return $ PieceInfo (pieceLen s) (HashList (BL.toStrict hashes))
 
+-- | Verify specific piece using infodict hash list.
+verifyPiece :: Storage -> PieceInfo -> PieceIx -> IO Bool
+verifyPiece s pinfo pix = checkPieceLazy pinfo <$> readPiece pix s
+
 -- | Verify storage.
 --
 --   Throws 'InvalidSize' if piece info size do not match with storage
@@ -174,11 +180,11 @@ genPieceInfo s = do
 getBitfield :: Storage -> PieceInfo -> IO Bitfield
 getBitfield s @ Storage {..} pinfo @ PieceInfo {..}
   | pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength)
-  | otherwise = M.foldM verifyPiece (BF.haveNone total) [0..total - 1]
+  | otherwise = M.foldM checkPiece (BF.haveNone total) [0..total - 1]
   where
     total = totalPieces s
 
-    verifyPiece :: Bitfield -> PieceIx -> IO Bitfield
-    verifyPiece bf pix = do
-      valid <- checkPieceLazy pinfo <$> readPiece pix s
+    checkPiece :: Bitfield -> PieceIx -> IO Bitfield
+    checkPiece bf pix = do
+      valid <- verifyPiece s pinfo pix
       return $ if valid then BF.insert pix bf else bf