diff options
Diffstat (limited to 'bittorrent/src/System')
-rw-r--r-- | bittorrent/src/System/Torrent/FileMap.hs | 163 | ||||
-rw-r--r-- | bittorrent/src/System/Torrent/Storage.hs | 221 | ||||
-rw-r--r-- | bittorrent/src/System/Torrent/Tree.hs | 83 |
3 files changed, 0 insertions, 467 deletions
diff --git a/bittorrent/src/System/Torrent/FileMap.hs b/bittorrent/src/System/Torrent/FileMap.hs deleted file mode 100644 index 38c475e8..00000000 --- a/bittorrent/src/System/Torrent/FileMap.hs +++ /dev/null | |||
@@ -1,163 +0,0 @@ | |||
1 | {-# LANGUAGE RecordWildCards #-} | ||
2 | {-# LANGUAGE ViewPatterns #-} | ||
3 | {-# OPTIONS -fno-warn-orphans #-} | ||
4 | module System.Torrent.FileMap | ||
5 | ( FileMap | ||
6 | |||
7 | -- * Construction | ||
8 | , Mode (..) | ||
9 | , def | ||
10 | , mmapFiles | ||
11 | , unmapFiles | ||
12 | |||
13 | -- * Query | ||
14 | , System.Torrent.FileMap.size | ||
15 | |||
16 | -- * Modification | ||
17 | , readBytes | ||
18 | , writeBytes | ||
19 | , unsafeReadBytes | ||
20 | |||
21 | -- * Unsafe conversions | ||
22 | , fromLazyByteString | ||
23 | , toLazyByteString | ||
24 | ) where | ||
25 | |||
26 | import Control.Applicative | ||
27 | import Control.Monad as L | ||
28 | import Data.ByteString as BS | ||
29 | import Data.ByteString.Internal as BS | ||
30 | import Data.ByteString.Lazy as BL | ||
31 | import Data.ByteString.Lazy.Internal as BL | ||
32 | import Data.Default | ||
33 | import Data.Vector as V -- TODO use unboxed vector | ||
34 | import Foreign | ||
35 | import System.IO.MMap | ||
36 | |||
37 | import Data.Torrent | ||
38 | |||
39 | |||
40 | data FileEntry = FileEntry | ||
41 | { filePosition :: {-# UNPACK #-} !FileOffset | ||
42 | , fileBytes :: {-# UNPACK #-} !BS.ByteString -- XXX: mutable buffer (see 'writeBytes'). | ||
43 | } deriving (Show, Eq) | ||
44 | |||
45 | type FileMap = Vector FileEntry | ||
46 | |||
47 | instance Default Mode where | ||
48 | def = ReadWriteEx | ||
49 | |||
50 | mmapFiles :: Mode -> FileLayout FileSize -> IO FileMap | ||
51 | mmapFiles mode layout = V.fromList <$> L.mapM mkEntry (accumPositions layout) | ||
52 | where | ||
53 | mkEntry (path, (pos, expectedSize)) = do | ||
54 | let esize = fromIntegral expectedSize -- FIXME does this safe? | ||
55 | (fptr, moff, msize) <- mmapFileForeignPtr path mode $ Just (0, esize) | ||
56 | if msize /= esize | ||
57 | then error "mmapFiles" -- TODO unmap mapped files on exception | ||
58 | else return $ FileEntry pos (PS fptr moff msize) | ||
59 | |||
60 | unmapFiles :: FileMap -> IO () | ||
61 | unmapFiles = V.mapM_ unmapEntry | ||
62 | where | ||
63 | unmapEntry (FileEntry _ (PS fptr _ _)) = finalizeForeignPtr fptr | ||
64 | |||
65 | -- Unsafe: FileMap 'writeBytes' will modify supplied bytestrings in place. | ||
66 | fromLazyByteString :: BL.ByteString -> FileMap | ||
67 | fromLazyByteString lbs = V.unfoldr f (0, lbs) | ||
68 | where | ||
69 | f (_, Empty ) = Nothing | ||
70 | f (pos, Chunk x xs) = Just (FileEntry pos x, ((pos + chunkSize), xs)) | ||
71 | where chunkSize = fromIntegral $ BS.length x | ||
72 | |||
73 | -- | /O(n)/. | ||
74 | -- | ||
75 | -- Unsafe: mutable buffers are returned without copy. | ||
76 | toLazyByteString :: FileMap -> BL.ByteString | ||
77 | toLazyByteString = V.foldr f Empty | ||
78 | where | ||
79 | f FileEntry {..} bs = Chunk fileBytes bs | ||
80 | |||
81 | -- | /O(1)/. | ||
82 | size :: FileMap -> FileOffset | ||
83 | size m | ||
84 | | V.null m = 0 | ||
85 | | FileEntry {..} <- V.unsafeLast m | ||
86 | = filePosition + fromIntegral (BS.length fileBytes) | ||
87 | |||
88 | -- | Find the file number for a particular byte offset within a torrent. | ||
89 | bsearch :: FileOffset -> FileMap -> Maybe Int | ||
90 | bsearch x m | ||
91 | | V.null m = Nothing | ||
92 | | otherwise = branch (V.length m `div` 2) | ||
93 | where | ||
94 | branch c @ ((m !) -> FileEntry {..}) | ||
95 | | x < filePosition = bsearch x (V.take c m) | ||
96 | | x >= filePosition + fileSize = do | ||
97 | ix <- bsearch x (V.drop (succ c) m) | ||
98 | return $ succ c + ix | ||
99 | | otherwise = Just c | ||
100 | where | ||
101 | fileSize = fromIntegral (BS.length fileBytes) | ||
102 | |||
103 | -- | /O(log n)/. | ||
104 | drop :: FileOffset -> FileMap -> (FileSize, FileMap) | ||
105 | drop off m | ||
106 | | Just ix <- bsearch off m | ||
107 | , FileEntry {..} <- m ! ix = (off - filePosition, V.drop ix m) | ||
108 | | otherwise = (0 , V.empty) | ||
109 | |||
110 | -- | /O(log n)/. | ||
111 | take :: FileSize -> FileMap -> (FileMap, FileSize) | ||
112 | take len m | ||
113 | | len >= s = (m , 0) | ||
114 | | Just ix <- bsearch (pred len) m = let m' = V.take (succ ix) m | ||
115 | in (m', System.Torrent.FileMap.size m' - len) | ||
116 | | otherwise = (V.empty , 0) | ||
117 | where | ||
118 | s = System.Torrent.FileMap.size m | ||
119 | |||
120 | -- | /O(log n + m)/. Do not use this function with 'unmapFiles'. | ||
121 | -- | ||
122 | -- The returned bytestring points directly into an area memory mapped from a | ||
123 | -- file. | ||
124 | unsafeReadBytes :: FileOffset -> FileSize -> FileMap -> BL.ByteString | ||
125 | unsafeReadBytes off s m | ||
126 | | (l , m') <- System.Torrent.FileMap.drop off m | ||
127 | , (m'', _ ) <- System.Torrent.FileMap.take (off + s) m' | ||
128 | = BL.take (fromIntegral s) $ BL.drop (fromIntegral l) $ toLazyByteString m'' | ||
129 | |||
130 | -- The returned bytestring is copied and safe to use after the file is | ||
131 | -- unmapped. | ||
132 | readBytes :: FileOffset -> FileSize -> FileMap -> IO BL.ByteString | ||
133 | readBytes off s m = do | ||
134 | let bs_copy = BL.copy $ unsafeReadBytes off s m | ||
135 | forceLBS bs_copy | ||
136 | return bs_copy | ||
137 | where | ||
138 | forceLBS Empty = return () | ||
139 | forceLBS (Chunk _ x) = forceLBS x | ||
140 | |||
141 | -- UNSAFE: Uses the first byte string as a pointer to mutable data and writes | ||
142 | -- the contents of the second bytestring there. | ||
143 | bscpy :: BL.ByteString -> BL.ByteString -> IO () | ||
144 | bscpy (PS _ _ 0 `Chunk` dest_rest) src = bscpy dest_rest src | ||
145 | bscpy dest (PS _ _ 0 `Chunk` src_rest) = bscpy dest src_rest | ||
146 | bscpy (PS dest_fptr dest_off dest_size `Chunk` dest_rest) | ||
147 | (PS src_fptr src_off src_size `Chunk` src_rest) | ||
148 | = do let csize = min dest_size src_size | ||
149 | withForeignPtr dest_fptr $ \dest_ptr -> | ||
150 | withForeignPtr src_fptr $ \src_ptr -> | ||
151 | memcpy (dest_ptr `advancePtr` dest_off) | ||
152 | (src_ptr `advancePtr` src_off) | ||
153 | (fromIntegral csize) -- TODO memmove? | ||
154 | bscpy (PS dest_fptr (dest_off + csize) (dest_size - csize) `Chunk` dest_rest) | ||
155 | (PS src_fptr (src_off + csize) (src_size - csize) `Chunk` src_rest) | ||
156 | bscpy _ _ = return () | ||
157 | |||
158 | -- UNSAFE: Mutates bytestring contents within the provided FileMap. | ||
159 | writeBytes :: FileOffset -> BL.ByteString -> FileMap -> IO () | ||
160 | writeBytes off lbs m = bscpy dest src | ||
161 | where | ||
162 | src = BL.take (fromIntegral (BL.length dest)) lbs | ||
163 | dest = unsafeReadBytes off (fromIntegral (BL.length lbs)) m | ||
diff --git a/bittorrent/src/System/Torrent/Storage.hs b/bittorrent/src/System/Torrent/Storage.hs deleted file mode 100644 index 1d77e55d..00000000 --- a/bittorrent/src/System/Torrent/Storage.hs +++ /dev/null | |||
@@ -1,221 +0,0 @@ | |||
1 | -- | | ||
2 | -- Copyright : (c) Sam Truzjan 2013 | ||
3 | -- License : BSD3 | ||
4 | -- Maintainer : pxqr.sta@gmail.com | ||
5 | -- Stability : experimental | ||
6 | -- Portability : portable | ||
7 | -- | ||
8 | -- This module implements mapping from single continious piece space | ||
9 | -- to file storage. Storage can be used in two modes: | ||
10 | -- | ||
11 | -- * As in memory storage - in this case we don't touch filesystem. | ||
12 | -- | ||
13 | -- * As ordinary mmaped file storage - when we need to store | ||
14 | -- data in the filesystem. | ||
15 | -- | ||
16 | {-# LANGUAGE RecordWildCards #-} | ||
17 | {-# LANGUAGE DeriveDataTypeable #-} | ||
18 | module System.Torrent.Storage | ||
19 | ( -- * Storage | ||
20 | Storage | ||
21 | , StorageFailure (..) | ||
22 | |||
23 | -- * Construction | ||
24 | , Mode (..) | ||
25 | , def | ||
26 | , open | ||
27 | , openInfoDict | ||
28 | , close | ||
29 | , withStorage | ||
30 | |||
31 | -- * Query | ||
32 | , totalPieces | ||
33 | , verifyPiece | ||
34 | , genPieceInfo | ||
35 | , getBitfield | ||
36 | |||
37 | -- * Modification | ||
38 | , writePiece | ||
39 | , readPiece | ||
40 | , hintRead | ||
41 | , unsafeReadPiece | ||
42 | |||
43 | -- * Streaming | ||
44 | , sourceStorage | ||
45 | , sinkStorage | ||
46 | ) where | ||
47 | |||
48 | import Control.Applicative | ||
49 | import Control.Exception | ||
50 | import Control.Monad as M | ||
51 | import Control.Monad.Trans | ||
52 | import Data.ByteString.Lazy as BL | ||
53 | import Data.Conduit as C | ||
54 | import Data.Conduit.Binary as C | ||
55 | import Data.Conduit.List as C | ||
56 | import Data.Typeable | ||
57 | |||
58 | import Data.Torrent | ||
59 | import Network.BitTorrent.Exchange.Bitfield as BF | ||
60 | import System.Torrent.FileMap as FM | ||
61 | |||
62 | |||
63 | -- | Some storage operations may throw an exception if misused. | ||
64 | data StorageFailure | ||
65 | -- | Occurs on a write operation if the storage has been opened | ||
66 | -- using 'ReadOnly' mode. | ||
67 | = StorageIsRO | ||
68 | |||
69 | -- | Piece index is out of bounds. | ||
70 | | InvalidIndex PieceIx | ||
71 | |||
72 | -- | Piece size do not match with one passed to the 'open' | ||
73 | -- function. | ||
74 | | InvalidSize PieceSize | ||
75 | deriving (Show, Eq, Typeable) | ||
76 | |||
77 | instance Exception StorageFailure | ||
78 | |||
79 | -- | Pieces store. | ||
80 | data Storage = Storage | ||
81 | { mode :: !Mode | ||
82 | , pieceLen :: {-# UNPACK #-} !PieceSize | ||
83 | , fileMap :: {-# UNPACK #-} !FileMap | ||
84 | } | ||
85 | |||
86 | -- | Map torrent files: | ||
87 | -- | ||
88 | -- * when torrent first created use 'ReadWriteEx' mode; | ||
89 | -- | ||
90 | -- * when seeding, validation 'ReadOnly' mode. | ||
91 | -- | ||
92 | open :: Mode -> PieceSize -> FileLayout FileSize -> IO Storage | ||
93 | open mode s l | ||
94 | | s <= 0 = throwIO (InvalidSize s) | ||
95 | | otherwise = Storage mode s <$> mmapFiles mode l | ||
96 | |||
97 | -- | Like 'open', but use 'InfoDict' file layout. | ||
98 | openInfoDict :: Mode -> FilePath -> InfoDict -> IO Storage | ||
99 | openInfoDict mode rootPath InfoDict {..} = | ||
100 | open mode (piPieceLength idPieceInfo) (flatLayout rootPath idLayoutInfo) | ||
101 | |||
102 | -- | Unmaps all files forcefully. It is recommended but not required. | ||
103 | close :: Storage -> IO () | ||
104 | close Storage {..} = unmapFiles fileMap | ||
105 | |||
106 | -- | Normally you need to use 'Control.Monad.Trans.Resource.allocate'. | ||
107 | withStorage :: Mode -> PieceSize -> FileLayout FileSize | ||
108 | -> (Storage -> IO ()) -> IO () | ||
109 | withStorage m s l = bracket (open m s l) close | ||
110 | |||
111 | -- TODO allocateStorage? | ||
112 | |||
113 | -- | Count of pieces in the storage. | ||
114 | totalPieces :: Storage -> PieceCount | ||
115 | totalPieces Storage {..} = FM.size fileMap `sizeInBase` pieceLen | ||
116 | |||
117 | isValidIx :: PieceIx -> Storage -> Bool | ||
118 | isValidIx i s = 0 <= i && i < totalPieces s | ||
119 | |||
120 | -- | Put piece data at the piece index by overwriting existing | ||
121 | -- data. | ||
122 | -- | ||
123 | -- This operation may throw 'StorageFailure'. | ||
124 | -- | ||
125 | writePiece :: Piece BL.ByteString -> Storage -> IO () | ||
126 | writePiece p @ Piece {..} s @ Storage {..} | ||
127 | | mode == ReadOnly = throwIO StorageIsRO | ||
128 | | isNotValidIx pieceIndex = throwIO (InvalidIndex pieceIndex) | ||
129 | | isNotValidSize pieceIndex (pieceSize p) | ||
130 | = throwIO (InvalidSize (pieceSize p)) | ||
131 | | otherwise = writeBytes offset pieceData fileMap | ||
132 | where | ||
133 | isNotValidSize pix psize | ||
134 | | succ pix == pcount = psize /= lastPieceLen -- last piece may be shorter | ||
135 | | otherwise = psize /= pieceLen | ||
136 | where | ||
137 | lastPieceLen = fromIntegral (FM.size fileMap `rem` fromIntegral pieceLen) | ||
138 | {-# INLINE isNotValidSize #-} | ||
139 | |||
140 | isNotValidIx i = i < 0 || i >= pcount | ||
141 | {-# INLINE isNotValidIx #-} | ||
142 | |||
143 | pcount = totalPieces s | ||
144 | offset = fromIntegral pieceIndex * fromIntegral pieceLen | ||
145 | |||
146 | -- | Read specific piece from storage. | ||
147 | -- | ||
148 | -- This operation may throw 'StorageFailure'. | ||
149 | -- | ||
150 | readPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString) | ||
151 | readPiece pix s @ Storage {..} | ||
152 | | not (isValidIx pix s) = throwIO (InvalidIndex pix) | ||
153 | | otherwise = Piece pix <$> readBytes offset sz fileMap | ||
154 | where | ||
155 | offset = fromIntegral pix * fromIntegral pieceLen | ||
156 | sz = fromIntegral pieceLen | ||
157 | |||
158 | -- | Hint about the coming 'readPiece'. Ignores invalid indexes, for e.g.: | ||
159 | -- | ||
160 | -- @forall s. hindRead (-1) s == return ()@ | ||
161 | -- | ||
162 | hintRead :: PieceIx -> Storage -> IO () | ||
163 | hintRead _pix Storage {..} = return () | ||
164 | |||
165 | -- | Zero-copy version of readPiece. Can be used only with 'ReadOnly' | ||
166 | -- storages. | ||
167 | unsafeReadPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString) | ||
168 | unsafeReadPiece pix s @ Storage {..} | ||
169 | | not (isValidIx pix s) = throwIO (InvalidIndex pix) | ||
170 | | otherwise = return $ Piece pix (unsafeReadBytes offset sz fileMap) | ||
171 | where | ||
172 | offset = fromIntegral pix * fromIntegral pieceLen | ||
173 | sz = fromIntegral pieceLen | ||
174 | |||
175 | -- | Stream storage pieces from first to the last. | ||
176 | sourceStorage :: Storage -> Source IO (Piece BL.ByteString) | ||
177 | sourceStorage s = go 0 | ||
178 | where | ||
179 | go pix | ||
180 | | pix < totalPieces s = do | ||
181 | piece <- liftIO $ readPiece pix s | ||
182 | liftIO $ hintRead (succ pix) s | ||
183 | yield piece | ||
184 | go (succ pix) | ||
185 | | otherwise = return () | ||
186 | |||
187 | -- | Write stream of pieces to the storage. Fail if storage is 'ReadOnly'. | ||
188 | sinkStorage :: Storage -> Sink (Piece BL.ByteString) IO () | ||
189 | sinkStorage s = do | ||
190 | awaitForever $ \ piece -> | ||
191 | liftIO $ writePiece piece s | ||
192 | |||
193 | -- | This function can be used to generate 'InfoDict' from a set of | ||
194 | -- opened files. | ||
195 | genPieceInfo :: Storage -> IO PieceInfo | ||
196 | genPieceInfo s = do | ||
197 | hashes <- sourceStorage s $= C.map hashPiece $$ C.sinkLbs | ||
198 | return $ PieceInfo (pieceLen s) (HashList (BL.toStrict hashes)) | ||
199 | |||
200 | -- | Verify specific piece using infodict hash list. | ||
201 | verifyPiece :: Storage -> PieceInfo -> PieceIx -> IO Bool | ||
202 | verifyPiece s pinfo pix = do | ||
203 | piece <- unsafeReadPiece pix s | ||
204 | return $! checkPieceLazy pinfo piece | ||
205 | |||
206 | -- | Verify storage. | ||
207 | -- | ||
208 | -- Throws 'InvalidSize' if piece info size do not match with storage | ||
209 | -- piece size. | ||
210 | -- | ||
211 | getBitfield :: Storage -> PieceInfo -> IO Bitfield | ||
212 | getBitfield s @ Storage {..} pinfo @ PieceInfo {..} | ||
213 | | pieceLen /= piPieceLength = throwIO (InvalidSize piPieceLength) | ||
214 | | otherwise = M.foldM checkPiece (BF.haveNone total) [0..total - 1] | ||
215 | where | ||
216 | total = totalPieces s | ||
217 | |||
218 | checkPiece :: Bitfield -> PieceIx -> IO Bitfield | ||
219 | checkPiece bf pix = do | ||
220 | valid <- verifyPiece s pinfo pix | ||
221 | return $ if valid then BF.insert pix bf else bf | ||
diff --git a/bittorrent/src/System/Torrent/Tree.hs b/bittorrent/src/System/Torrent/Tree.hs deleted file mode 100644 index 41cfb360..00000000 --- a/bittorrent/src/System/Torrent/Tree.hs +++ /dev/null | |||
@@ -1,83 +0,0 @@ | |||
1 | -- | | ||
2 | -- Copyright : (c) Sam Truzjan 2013 | ||
3 | -- License : BSD3 | ||
4 | -- Maintainer : pxqr.sta@gmail.com | ||
5 | -- Stability : experimental | ||
6 | -- Portability : portable | ||
7 | -- | ||
8 | -- Directory tree can be used to easily manipulate file layout info. | ||
9 | -- | ||
10 | {-# LANGUAGE FlexibleInstances #-} | ||
11 | {-# LANGUAGE TemplateHaskell #-} | ||
12 | {-# LANGUAGE DeriveDataTypeable #-} | ||
13 | module System.Torrent.Tree | ||
14 | ( -- * Directory tree | ||
15 | DirTree (..) | ||
16 | |||
17 | -- * Construction | ||
18 | , build | ||
19 | |||
20 | -- * Query | ||
21 | , System.Torrent.Tree.lookup | ||
22 | , lookupDir | ||
23 | , fileNumber | ||
24 | , dirNumber | ||
25 | ) where | ||
26 | |||
27 | import Data.ByteString as BS | ||
28 | import Data.ByteString.Char8 as BC | ||
29 | import Data.Foldable | ||
30 | import Data.List as L | ||
31 | import Data.Map as M | ||
32 | import Data.Monoid | ||
33 | |||
34 | import Data.Torrent | ||
35 | |||
36 | |||
37 | -- | 'DirTree' is more convenient form of 'LayoutInfo'. | ||
38 | data DirTree a = Dir { children :: Map ByteString (DirTree a) } | ||
39 | | File { node :: FileInfo a } | ||
40 | deriving Show | ||
41 | |||
42 | -- | Build directory tree from a list of files. | ||
43 | build :: LayoutInfo -> DirTree () | ||
44 | build SingleFile {liFile = FileInfo {..}} = Dir | ||
45 | { children = M.singleton fiName (File fi) } | ||
46 | where | ||
47 | fi = FileInfo fiLength fiMD5Sum () | ||
48 | build MultiFile {..} = Dir $ M.singleton liDirName files | ||
49 | where | ||
50 | files = Dir $ M.fromList $ L.map mkFileEntry liFiles | ||
51 | mkFileEntry FileInfo {..} = (L.head fiName, ent) -- TODO FIXME | ||
52 | where | ||
53 | ent = File $ FileInfo fiLength fiMD5Sum () | ||
54 | |||
55 | --decompress :: DirTree () -> [FileInfo ()] | ||
56 | --decompress = undefined | ||
57 | |||
58 | -- TODO pretty print | ||
59 | |||
60 | -- | Lookup file by path. | ||
61 | lookup :: [FilePath] -> DirTree a -> Maybe (DirTree a) | ||
62 | lookup [] t = Just t | ||
63 | lookup (p : ps) (Dir m) | Just subTree <- M.lookup (BC.pack p) m | ||
64 | = System.Torrent.Tree.lookup ps subTree | ||
65 | lookup _ _ = Nothing | ||
66 | |||
67 | -- | Lookup directory by path. | ||
68 | lookupDir :: [FilePath] -> DirTree a -> Maybe [(ByteString, DirTree a)] | ||
69 | lookupDir ps d = do | ||
70 | subTree <- System.Torrent.Tree.lookup ps d | ||
71 | case subTree of | ||
72 | File _ -> Nothing | ||
73 | Dir es -> Just $ M.toList es | ||
74 | |||
75 | -- | Get total count of files in directory and subdirectories. | ||
76 | fileNumber :: DirTree a -> Sum Int | ||
77 | fileNumber File {..} = Sum 1 | ||
78 | fileNumber Dir {..} = foldMap fileNumber children | ||
79 | |||
80 | -- | Get total count of directories in the directory and subdirectories. | ||
81 | dirNumber :: DirTree a -> Sum Int | ||
82 | dirNumber File {..} = Sum 0 | ||
83 | dirNumber Dir {..} = Sum 1 <> foldMap dirNumber children | ||