src/FuncTorrent/FileSystem.hs

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