src/FuncTorrent/FileSystem.hs

   1 {-
   2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   3  -
   4  - This file is part of FuncTorrent.
   5  -
   6  - FuncTorrent is free software; you can redistribute it and/or modify
   7  - it under the terms of the GNU General Public License as published by
   8  - the Free Software Foundation; either version 3 of the License, or
   9  - (at your option) any later version.
  10  -
  11  - FuncTorrent is distributed in the hope that it will be useful,
  12  - but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  - GNU General Public License for more details.
  15  -
  16  - You should have received a copy of the GNU General Public License
  17  - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
  18  -}
  19
  20 {-# LANGUAGE OverloadedStrings #-}
  21 {-# LANGUAGE FlexibleContexts #-}
  22
  23 module FuncTorrent.FileSystem
  24        (run,
  25         MsgChannel,
  26         createMsgChannel,
  27         writePieceToDisk,
  28         Piece(..),
  29         pieceMapFromFile,
  30         Stats(..),
  31         getStats
  32        )
  33        where
  34
  35 import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
  36 import           Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar)
  37 import           Control.Monad (forever)
  38 import           Control.Monad.State (StateT, liftIO, get, runStateT, modify)
  39 import qualified Data.ByteString as BS
  40 import           Data.Map (traverseWithKey, (!))
  41 import           System.IO (Handle, IOMode (ReadWriteMode), withFile)
  42 import           System.Directory (doesFileExist)
  43
  44 import           FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pieceNumToOffset)
  45 import           FuncTorrent.Utils (createDummyFile, readFileAtOffset, writeFileAtOffset, verifyHash)
  46
  47 type PieceNum = Integer
  48 data Piece = Piece PieceNum BS.ByteString
  49
  50 data Msg = ReadPiece PieceNum Integer (MVar Piece)
  51          | WritePiece Piece
  52          | VerifyPiece PieceNum (MVar Bool)
  53          | GetStats (MVar Stats)
  54
  55 type MsgChannel = Chan Msg
  56
  57 data Stats = Stats { bytesRead :: Integer
  58                    , bytesWritten :: Integer
  59                    }
  60
  61 createMsgChannel :: IO (Chan Msg)
  62 createMsgChannel = newChan
  63
  64 run :: PieceMap -> MsgChannel -> Handle -> IO ()
  65 run pieceMap c handle = forever $ do
  66   _ <- runStateT (run' pieceMap c handle) initialStats
  67   return ()
  68     where initialStats = Stats { bytesRead = 0
  69                                , bytesWritten = 0 }
  70
  71 run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
  72 run' pieceMap c handle = do
  73   stats <- get
  74   msg <- liftIO recvMsg
  75   liftIO $ sendResponse msg stats
  76   updateStats msg
  77     where
  78       recvMsg = readChan c
  79       sendResponse msg stats =
  80         case msg of
  81           ReadPiece n len' var -> do
  82             bs <- readPiece n len'
  83             putMVar var (Piece n bs)
  84           WritePiece (Piece n bs) ->
  85             writePiece n bs
  86           VerifyPiece n var -> do
  87             isHashValid <- verifyPiece n
  88             putMVar var isHashValid
  89           GetStats var ->
  90             putMVar var stats
  91       readPiece n len' = do
  92         let offset = pieceNumToOffset pieceMap n
  93         readFileAtOffset handle offset len'
  94       writePiece n piece = do
  95         let offset = pieceNumToOffset pieceMap n
  96         writeFileAtOffset handle offset piece
  97       verifyPiece n = do
  98         let offset = pieceNumToOffset pieceMap n
  99             hash'  = hash (pieceMap ! n)
 100             len'   = len (pieceMap ! n)
 101         bs' <- readFileAtOffset handle offset len'
 102         return $ verifyHash bs' hash'
 103       updateStats (ReadPiece _ l _) =
 104         modify (\st -> st {bytesRead = bytesRead st + l})
 105       updateStats (WritePiece (Piece _ bs)) =
 106         modify (\st -> st {bytesWritten = bytesWritten st + fromIntegral (BS.length bs)})
 107       updateStats _ = modify id
 108
 109 pieceMapFromFile :: FilePath -> Integer -> PieceMap -> IO PieceMap
 110 pieceMapFromFile filePath fileLen pieceMap = do
 111   dfe <- doesFileExist filePath
 112   if dfe
 113     then traverseWithKey f pieceMap
 114     else createDummyFile filePath (fromIntegral fileLen) >> return pieceMap
 115   where
 116     f k v = do
 117       let offset = pieceNumToOffset pieceMap k
 118       isHashValid <- flip verifyHash (hash v) <$> withFile filePath ReadWriteMode (\handle -> readFileAtOffset handle offset (len v))
 119       if isHashValid
 120         then return $ v { dlstate = Have }
 121         else return v
 122
 123 writePieceToDisk :: MsgChannel -> PieceNum -> BS.ByteString -> IO ()
 124 writePieceToDisk c pieceNum bs =
 125   writeChan c $ WritePiece (Piece pieceNum bs)
 126
 127 getStats :: MsgChannel -> IO (MVar Stats)
 128 getStats c = do
 129   v <- newEmptyMVar
 130   writeChan c $ GetStats v
 131   return v