]> git.rkrishnan.org Git - functorrent.git/blob - src/FuncTorrent/FileSystem.hs
tests: add license boilerplate
[functorrent.git] / 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