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         Stats(..),
  11         getStats
  12        )
  13        where
  14
  15 import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
  16 import           Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar)
  17 import           Control.Monad (forever)
  18 import           Control.Monad.State (StateT, liftIO, get, runStateT, modify)
  19 import qualified Data.ByteString as BS
  20 import           Data.Map (traverseWithKey, (!))
  21 import           System.IO (Handle, IOMode (ReadWriteMode), withFile)
  22 import           System.Directory (doesFileExist)
  23
  24 import           FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pieceNumToOffset)
  25 import           FuncTorrent.Utils (createDummyFile, readFileAtOffset, writeFileAtOffset, verifyHash)
  26
  27 type PieceNum = Integer
  28 data Piece = Piece PieceNum BS.ByteString
  29
  30 data Msg = ReadPiece PieceNum Integer (MVar Piece)
  31          | WritePiece Piece
  32          | VerifyPiece PieceNum (MVar Bool)
  33          | GetStats (MVar Stats)
  34
  35 type MsgChannel = Chan Msg
  36
  37 data Stats = Stats { bytesRead :: Integer
  38                    , bytesWritten :: Integer
  39                    }
  40
  41 createMsgChannel :: IO (Chan Msg)
  42 createMsgChannel = newChan
  43
  44 run :: PieceMap -> MsgChannel -> Handle -> IO ()
  45 run pieceMap c handle = forever $ do
  46   _ <- runStateT (run' pieceMap c handle) initialStats
  47   return ()
  48     where initialStats = Stats { bytesRead = 0
  49                                , bytesWritten = 0 }
  50
  51 run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
  52 run' pieceMap c handle = do
  53   stats <- get
  54   msg <- liftIO recvMsg
  55   liftIO $ sendResponse msg stats
  56   updateStats msg
  57     where
  58       recvMsg = readChan c
  59       sendResponse msg stats =
  60         case msg of
  61           ReadPiece n len' var -> do
  62             bs <- readPiece n len'
  63             putMVar var (Piece n bs)
  64           WritePiece (Piece n bs) ->
  65             writePiece n bs
  66           VerifyPiece n var -> do
  67             isHashValid <- verifyPiece n
  68             putMVar var isHashValid
  69           GetStats var -> do
  70             putMVar var stats
  71       readPiece n len' = do
  72         let offset = pieceNumToOffset pieceMap n
  73         readFileAtOffset handle offset len'
  74       writePiece n piece = do
  75         let offset = pieceNumToOffset pieceMap n
  76         writeFileAtOffset handle offset piece
  77       verifyPiece n = do
  78         let offset = pieceNumToOffset pieceMap n
  79             hash'  = hash (pieceMap ! n)
  80             len'   = len (pieceMap ! n)
  81         bs' <- readFileAtOffset handle offset len'
  82         return $ verifyHash bs' hash'
  83       updateStats (ReadPiece _ l _) =
  84         modify (\st -> st {bytesRead = bytesRead st + l})
  85       updateStats (WritePiece (Piece _ bs)) =
  86         modify (\st -> st {bytesWritten = bytesWritten st + fromIntegral (BS.length bs)})
  87       updateStats _ = modify id
  88
  89 pieceMapFromFile :: FilePath -> Integer -> PieceMap -> IO PieceMap
  90 pieceMapFromFile filePath fileLen pieceMap = do
  91   dfe <- doesFileExist filePath
  92   if dfe
  93     then traverseWithKey f pieceMap
  94     else createDummyFile filePath (fromIntegral fileLen) >> return pieceMap
  95   where
  96     f k v = do
  97       let offset = pieceNumToOffset pieceMap k
  98       isHashValid <- flip verifyHash (hash v) <$> withFile filePath ReadWriteMode (\handle -> readFileAtOffset handle offset (len v))
  99       if isHashValid
 100         then return $ v { dlstate = Have }
 101         else return v
 102
 103 writePieceToDisk :: MsgChannel -> PieceNum -> BS.ByteString -> IO ()
 104 writePieceToDisk c pieceNum bs =
 105   writeChan c $ WritePiece (Piece pieceNum bs)
 106
 107 getStats :: MsgChannel -> IO (MVar Stats)
 108 getStats c = do
 109   v <- newEmptyMVar
 110   writeChan c $ GetStats v
 111   return v