src/FuncTorrent/FileSystem.hs

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