src/FuncTorrent/FileSystem.hs

   1 module FuncTorrent.FileSystem
   2        (startThread,
   3         MsgChannel,
   4         createMsgChannel,
   5         Msg(..),
   6         Piece(..),
   7         pieceMapFromFile
   8        )
   9        where
  10
  11 import           Control.Concurrent (ThreadId, forkIO)
  12 import           Control.Concurrent.Chan (Chan, newChan, readChan)
  13 import           Control.Concurrent.MVar (MVar, putMVar)
  14 import           Control.Monad (forever)
  15 import           Data.Map (traverseWithKey)
  16
  17 import qualified Data.ByteString as BS
  18 import           Data.Map ((!))
  19 import           System.IO (Handle, IOMode (ReadWriteMode), withFile)
  20 import           System.Directory (doesFileExist)
  21
  22 import           FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pieceNumToOffset)
  23 import           FuncTorrent.Utils (createDummyFile, readFileAtOffset, writeFileAtOffset, verifyHash)
  24
  25 type PieceNum = Integer
  26 data Piece = Piece PieceNum BS.ByteString
  27
  28 data Msg = ReadPiece PieceNum Integer (MVar Piece)
  29          | WritePiece Piece
  30          | VerifyPiece PieceNum (MVar Bool)
  31
  32 type MsgChannel = Chan Msg
  33
  34 createMsgChannel :: IO (Chan Msg)
  35 createMsgChannel = newChan
  36
  37 startThread :: PieceMap -> MsgChannel -> Handle -> IO ThreadId
  38 startThread pieceMap c handle = do
  39   forkIO $ forever $ recvMsg >>= sendResponse
  40   where
  41     recvMsg = readChan c
  42     sendResponse msg =
  43       case msg of
  44         ReadPiece n len' var -> do
  45           bs <- readPiece n len'
  46           putMVar var (Piece n bs)
  47         WritePiece (Piece n bs) -> do
  48           writePiece n bs
  49         VerifyPiece n var -> do
  50           isHashValid <- verifyPiece n
  51           putMVar var isHashValid
  52     readPiece n len' = do
  53       let offset = pieceNumToOffset pieceMap n
  54       readFileAtOffset handle offset len'
  55     writePiece n piece = do
  56       let offset = pieceNumToOffset pieceMap n
  57       writeFileAtOffset handle offset piece
  58     verifyPiece n = do
  59       let offset = pieceNumToOffset pieceMap n
  60           hash'  = hash (pieceMap ! n)
  61           len'   = len (pieceMap ! n)
  62       bs' <- readFileAtOffset handle offset len'
  63       return $ verifyHash bs' hash'
  64
  65 pieceMapFromFile :: FilePath -> Integer -> PieceMap -> IO PieceMap
  66 pieceMapFromFile filePath fileLen pieceMap = do
  67   dfe <- doesFileExist filePath
  68   if dfe
  69     then traverseWithKey f pieceMap
  70     else createDummyFile filePath (fromIntegral fileLen) >> return pieceMap
  71   where
  72     f k v = do
  73       let offset = pieceNumToOffset pieceMap k
  74       isHashValid <- flip verifyHash (hash v) <$> withFile filePath ReadWriteMode (\handle -> readFileAtOffset handle offset (len v))
  75       if isHashValid
  76         then return $ v { dlstate = Have }
  77         else return v