1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
|
-- |
-- Copyright : (c) Sam Truzjan 2013
-- License : BSD3
-- Maintainer : pxqr.sta@gmail.com
-- Stability : experimental
-- Portability : portable
--
-- This module implements mapping from single continious piece space
-- to file storage. Storage can be used in two modes:
--
-- * As in memory storage - in this case we don't touch filesystem.
--
-- * As ordinary mmaped file storage - when we need to store
-- data in the filesystem.
--
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveDataTypeable #-}
module System.Torrent.Storage
( -- * Storage
Storage
, StorageFailure (..)
-- * Construction
, Mode (..)
, def
, open
, close
, withStorage
-- * Query
, totalPieces
, verifyPiece
, genPieceInfo
, getBitfield
-- * Modification
, writePiece
, readPiece
, hintRead
, unsafeReadPiece
-- * Streaming
, sourceStorage
, sinkStorage
) where
import Control.Applicative
import Control.Exception
import Control.Monad as M
import Control.Monad.Trans
import Data.ByteString.Lazy as BL
import Data.Conduit as C
import Data.Conduit.Binary as C
import Data.Conduit.List as C
import Data.Typeable
import Data.Torrent.Bitfield as BF
import Data.Torrent.Layout
import Data.Torrent.Piece
import System.Torrent.FileMap as FM
data StorageFailure
-- | Occurs on a write operation if the storage has been opened
-- using 'ReadOnly' mode.
= StorageIsRO
-- | Piece index is out of bounds.
| InvalidIndex PieceIx
-- | Piece size do not match with one passed to the 'open'
-- function.
| InvalidSize PieceSize
deriving (Show, Eq, Typeable)
instance Exception StorageFailure
-- TODO validation
data Storage = Storage
{ mode :: !Mode
, pieceLen :: {-# UNPACK #-} !PieceSize
, fileMap :: {-# UNPACK #-} !FileMap
}
-- ResourceT ?
open :: Mode -> PieceSize -> FileLayout FileSize -> IO Storage
open mode s l = Storage mode s <$> mmapFiles mode l
close :: Storage -> IO ()
close Storage {..} = unmapFiles fileMap
withStorage :: Mode -> PieceSize -> FileLayout FileSize
-> (Storage -> IO ()) -> IO ()
withStorage m s l = bracket (open m s l) close
totalPieces :: Storage -> PieceCount
totalPieces Storage {..} = FM.size fileMap `sizeInBase` pieceLen
isValidIx :: PieceIx -> Storage -> Bool
isValidIx i s = 0 <= i && i < totalPieces s
writePiece :: Piece BL.ByteString -> Storage -> IO ()
writePiece p @ Piece {..} s @ Storage {..}
| mode == ReadOnly = throwIO StorageIsRO
| isNotValidIx pieceIndex = throwIO (InvalidIndex pieceIndex)
| isNotValidSize pieceIndex (pieceSize p)
= throwIO (InvalidSize (pieceSize p))
| otherwise = writeBytes offset pieceData fileMap
where
isNotValidSize pix psize
| succ pix == pcount = psize /= lastPieceLen -- last piece may be shorter
| otherwise = psize /= pieceLen
where
lastPieceLen = fromIntegral (FM.size fileMap `rem` fromIntegral pieceLen)
{-# INLINE isNotValidSize #-}
isNotValidIx i = i < 0 || i >= pcount
{-# INLINE isNotValidIx #-}
pcount = totalPieces s
offset = fromIntegral pieceIndex * fromIntegral pieceLen
readPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
readPiece pix s @ Storage {..}
| not (isValidIx pix s) = throwIO (InvalidIndex pix)
| otherwise = Piece pix <$> readBytes offset sz fileMap
where
offset = fromIntegral pix * fromIntegral pieceLen
sz = fromIntegral pieceLen
-- | Hint about the coming 'readPiece'. Ignores invalid indexes, for e.g.:
--
-- @forall s. hindRead (-1) s == return ()@
--
hintRead :: PieceIx -> Storage -> IO ()
hintRead _pix Storage {..} = return ()
-- | Zero-copy version of readPiece. Can be used only with 'ReadOnly'
-- storages.
unsafeReadPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
unsafeReadPiece pix s @ Storage {..}
| not (isValidIx pix s) = throwIO (InvalidIndex pix)
| otherwise = return $ Piece pix (unsafeReadBytes offset sz fileMap)
where
offset = fromIntegral pix * fromIntegral pieceLen
sz = fromIntegral pieceLen
-- | Stream storage pieces from first to the last.
sourceStorage :: Storage -> Source IO (Piece BL.ByteString)
sourceStorage s = go 0
where
go pix
| pix < totalPieces s = do
piece <- liftIO $ readPiece pix s
liftIO $ hintRead (succ pix) s
yield piece
go (succ pix)
| otherwise = return ()
-- | Write stream of pieces to the storage. Fail if storage is 'ReadOnly'.
sinkStorage :: Storage -> Sink (Piece BL.ByteString) IO ()
sinkStorage s = do
awaitForever $ \ piece ->
liftIO $ writePiece piece s
-- | This function can be used to generate 'InfoDict' from a set of
-- opened files.
genPieceInfo :: Storage -> IO PieceInfo
genPieceInfo s = do
hashes <- sourceStorage s $= C.map hashPiece $$ C.sinkLbs
return $ PieceInfo (pieceLen s) (HashList (BL.toStrict hashes))
-- | Verify specific piece using infodict hash list.
verifyPiece :: Storage -> PieceInfo -> PieceIx -> IO Bool
verifyPiece s pinfo pix = do
piece <- unsafeReadPiece pix s
return $! checkPieceLazy pinfo piece
-- | Verify storage.
--
-- Throws 'InvalidSize' if piece info size do not match with storage
-- piece size.
--
getBitfield :: Storage -> PieceInfo -> IO Bitfield
getBitfield s @ Storage {..} pinfo @ PieceInfo {..}
| pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength)
| otherwise = M.foldM checkPiece (BF.haveNone total) [0..total - 1]
where
total = totalPieces s
checkPiece :: Bitfield -> PieceIx -> IO Bitfield
checkPiece bf pix = do
valid <- verifyPiece s pinfo pix
return $ if valid then BF.insert pix bf else bf
|
