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
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
|
-- |
-- 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
, openInfoDict
, 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
import Network.BitTorrent.Exchange.Bitfield as BF
import System.Torrent.FileMap as FM
-- | Some storage operations may throw an exception if misused.
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
-- | Pieces store.
data Storage = Storage
{ mode :: !Mode
, pieceLen :: {-# UNPACK #-} !PieceSize
, fileMap :: {-# UNPACK #-} !FileMap
}
-- | Map torrent files:
--
-- * when torrent first created use 'ReadWriteEx' mode;
--
-- * when seeding, validation 'ReadOnly' mode.
--
open :: Mode -> PieceSize -> FileLayout FileSize -> IO Storage
open mode s l
| s <= 0 = throwIO (InvalidSize s)
| otherwise = Storage mode s <$> mmapFiles mode l
-- | Like 'open', but use 'InfoDict' file layout.
openInfoDict :: Mode -> FilePath -> InfoDict -> IO Storage
openInfoDict mode rootPath InfoDict {..} =
open mode (piPieceLength idPieceInfo) (flatLayout rootPath idLayoutInfo)
-- | Unmaps all files forcefully. It is recommended but not required.
close :: Storage -> IO ()
close Storage {..} = unmapFiles fileMap
-- | Normally you need to use 'Control.Monad.Trans.Resource.allocate'.
withStorage :: Mode -> PieceSize -> FileLayout FileSize
-> (Storage -> IO ()) -> IO ()
withStorage m s l = bracket (open m s l) close
-- TODO allocateStorage?
-- | Count of pieces in the storage.
totalPieces :: Storage -> PieceCount
totalPieces Storage {..} = FM.size fileMap `sizeInBase` pieceLen
isValidIx :: PieceIx -> Storage -> Bool
isValidIx i s = 0 <= i && i < totalPieces s
-- | Put piece data at the piece index by overwriting existing
-- data.
--
-- This operation may throw 'StorageFailure'.
--
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
-- | Read specific piece from storage.
--
-- This operation may throw 'StorageFailure'.
--
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
|
