src/FuncTorrent/FileSystem.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 module FuncTorrent.FileSystem
   4        (run,
   5         MsgChannel,
   6         createMsgChannel,
   7         writePieceToDisk,
   8         Piece(..),
   9         pieceMapFromFile
  10        )
  11        where
  12
  13 import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
  14 import           Control.Concurrent.MVar (MVar, putMVar)
  15 import           Control.Monad (forever)
  16 import           Control.Monad.State (StateT, liftIO, runStateT, modify)
  17 import qualified Data.ByteString as BS
  18 import           Data.Map (traverseWithKey, (!))
  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 data Stats = Stats { bytesRead :: Integer
  35                    , bytesWritten :: Integer
  36                    }
  37
  38 createMsgChannel :: IO (Chan Msg)
  39 createMsgChannel = newChan
  40
  41 run :: PieceMap -> MsgChannel -> Handle -> IO ()
  42 run pieceMap c handle = forever $ do
  43   _ <- runStateT (run' pieceMap c handle) initialStats
  44   return ()
  45     where initialStats = Stats { bytesRead = 0
  46                                , bytesWritten = 0 }
  47
  48 run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
  49 run' pieceMap c handle = do
  50   msg <- liftIO recvMsg
  51   liftIO $ sendResponse msg
  52   updateStats msg
  53   where
  54     recvMsg = readChan c
  55     sendResponse msg =
  56       case msg of
  57       ReadPiece n len' var -> do
  58         bs <- readPiece n len'
  59         putMVar var (Piece n bs)
  60       WritePiece (Piece n bs) ->
  61         writePiece n bs
  62       VerifyPiece n var -> do
  63         isHashValid <- verifyPiece n
  64         putMVar var isHashValid
  65     readPiece n len' = do
  66       let offset = pieceNumToOffset pieceMap n
  67       readFileAtOffset handle offset len'
  68     writePiece n piece = do
  69       let offset = pieceNumToOffset pieceMap n
  70       writeFileAtOffset handle offset piece
  71     verifyPiece n = do
  72       let offset = pieceNumToOffset pieceMap n
  73           hash'  = hash (pieceMap ! n)
  74           len'   = len (pieceMap ! n)
  75       bs' <- readFileAtOffset handle offset len'
  76       return $ verifyHash bs' hash'
  77     updateStats (ReadPiece _ l _) =
  78       modify (\st -> st {bytesRead = bytesRead st + l})
  79     updateStats (WritePiece (Piece _ bs)) =
  80       modify (\st -> st {bytesWritten = bytesWritten st + fromIntegral (BS.length bs)})
  81     updateStats _ = modify id
  82
  83 pieceMapFromFile :: FilePath -> Integer -> PieceMap -> IO PieceMap
  84 pieceMapFromFile filePath fileLen pieceMap = do
  85   dfe <- doesFileExist filePath
  86   if dfe
  87     then traverseWithKey f pieceMap
  88     else createDummyFile filePath (fromIntegral fileLen) >> return pieceMap
  89   where
  90     f k v = do
  91       let offset = pieceNumToOffset pieceMap k
  92       isHashValid <- flip verifyHash (hash v) <$> withFile filePath ReadWriteMode (\handle -> readFileAtOffset handle offset (len v))
  93       if isHashValid
  94         then return $ v { dlstate = Have }
  95         else return v
  96
  97 writePieceToDisk :: MsgChannel -> PieceNum -> BS.ByteString -> IO ()
  98 writePieceToDisk c pieceNum bs =
  99   writeChan c $ WritePiece (Piece pieceNum bs)
 100