1 | {-# LANGUAGE ForeignFunctionInterface, EmptyDataDecls #-} |
---|
2 | -- | |
---|
3 | -- Module: Codec.FEC |
---|
4 | -- Copyright: Adam Langley |
---|
5 | -- License: GPLv2+|TGPPLv1+ (see README.rst for details) |
---|
6 | -- |
---|
7 | -- Stability: experimental |
---|
8 | -- |
---|
9 | -- The module provides k of n encoding - a way to generate (n - k) secondary |
---|
10 | -- blocks of data from k primary blocks such that any k blocks (primary or |
---|
11 | -- secondary) are sufficient to regenerate all blocks. |
---|
12 | -- |
---|
13 | -- All blocks must be the same length and you need to keep track of which |
---|
14 | -- blocks you have in order to tell decode. By convention, the blocks are |
---|
15 | -- numbered 0..(n - 1) and blocks numbered < k are the primary blocks. |
---|
16 | |
---|
17 | module Codec.FEC ( |
---|
18 | FECParams |
---|
19 | , fec |
---|
20 | , encode |
---|
21 | , decode |
---|
22 | |
---|
23 | -- * Utility functions |
---|
24 | , secureDivide |
---|
25 | , secureCombine |
---|
26 | , enFEC |
---|
27 | , deFEC |
---|
28 | ) where |
---|
29 | |
---|
30 | import qualified Data.ByteString as B |
---|
31 | import qualified Data.ByteString.Unsafe as BU |
---|
32 | import qualified Data.ByteString.Internal as BI |
---|
33 | import Data.Word (Word8) |
---|
34 | import Data.Bits (xor) |
---|
35 | import Data.List (sortBy, partition, (\\), nub) |
---|
36 | import Foreign.Ptr |
---|
37 | import Foreign.Storable (sizeOf, poke) |
---|
38 | import Foreign.ForeignPtr |
---|
39 | import Foreign.C.Types |
---|
40 | import Foreign.Marshal.Alloc |
---|
41 | import Foreign.Marshal.Array (withArray, advancePtr) |
---|
42 | import System.IO (withFile, IOMode(..)) |
---|
43 | import System.IO.Unsafe (unsafePerformIO) |
---|
44 | |
---|
45 | data CFEC |
---|
46 | data FECParams = FECParams (ForeignPtr CFEC) Int Int |
---|
47 | |
---|
48 | instance Show FECParams where |
---|
49 | show (FECParams _ k n) = "FEC (" ++ show k ++ ", " ++ show n ++ ")" |
---|
50 | |
---|
51 | foreign import ccall unsafe "fec_new" _new :: CUInt -- ^ k |
---|
52 | -> CUInt -- ^ n |
---|
53 | -> IO (Ptr CFEC) |
---|
54 | foreign import ccall unsafe "&fec_free" _free :: FunPtr (Ptr CFEC -> IO ()) |
---|
55 | foreign import ccall unsafe "fec_encode" _encode :: Ptr CFEC |
---|
56 | -> Ptr (Ptr Word8) -- ^ primary blocks |
---|
57 | -> Ptr (Ptr Word8) -- ^ (output) secondary blocks |
---|
58 | -> Ptr CUInt -- ^ array of secondary block ids |
---|
59 | -> CSize -- ^ length of previous |
---|
60 | -> CSize -- ^ block length |
---|
61 | -> IO () |
---|
62 | foreign import ccall unsafe "fec_decode" _decode :: Ptr CFEC |
---|
63 | -> Ptr (Ptr Word8) -- ^ input blocks |
---|
64 | -> Ptr (Ptr Word8) -- ^ output blocks |
---|
65 | -> Ptr CUInt -- ^ array of input indexes |
---|
66 | -> CSize -- ^ block length |
---|
67 | -> IO () |
---|
68 | |
---|
69 | -- | Return true if the given @k@ and @n@ values are valid |
---|
70 | isValidConfig :: Int -> Int -> Bool |
---|
71 | isValidConfig k n |
---|
72 | | k >= n = False |
---|
73 | | k < 1 = False |
---|
74 | | n < 1 = False |
---|
75 | | n > 255 = False |
---|
76 | | otherwise = True |
---|
77 | |
---|
78 | -- | Return a FEC with the given parameters. |
---|
79 | fec :: Int -- ^ the number of primary blocks |
---|
80 | -> Int -- ^ the total number blocks, must be < 256 |
---|
81 | -> FECParams |
---|
82 | fec k n = |
---|
83 | if not (isValidConfig k n) |
---|
84 | then error $ "Invalid FEC parameters: " ++ show k ++ " " ++ show n |
---|
85 | else unsafePerformIO (do |
---|
86 | cfec <- _new (fromIntegral k) (fromIntegral n) |
---|
87 | params <- newForeignPtr _free cfec |
---|
88 | return $ FECParams params k n) |
---|
89 | |
---|
90 | -- | Create a C array of unsigned from an input array |
---|
91 | uintCArray :: [Int] -> ((Ptr CUInt) -> IO a) -> IO a |
---|
92 | uintCArray xs f = withArray (map fromIntegral xs) f |
---|
93 | |
---|
94 | -- | Convert a list of ByteStrings to an array of pointers to their data |
---|
95 | byteStringsToArray :: [B.ByteString] -> ((Ptr (Ptr Word8)) -> IO a) -> IO a |
---|
96 | byteStringsToArray inputs f = do |
---|
97 | let l = length inputs |
---|
98 | allocaBytes (l * sizeOf (undefined :: Ptr Word8)) (\array -> do |
---|
99 | let inner _ [] = f array |
---|
100 | inner array' (bs : bss) = BU.unsafeUseAsCString bs (\ptr -> do |
---|
101 | poke array' $ castPtr ptr |
---|
102 | inner (advancePtr array' 1) bss) |
---|
103 | inner array inputs) |
---|
104 | |
---|
105 | -- | Return True iff all the given ByteStrings are the same length |
---|
106 | allByteStringsSameLength :: [B.ByteString] -> Bool |
---|
107 | allByteStringsSameLength [] = True |
---|
108 | allByteStringsSameLength (bs : bss) = all ((==) (B.length bs)) $ map B.length bss |
---|
109 | |
---|
110 | -- | Run the given function with a pointer to an array of @n@ pointers to |
---|
111 | -- buffers of size @size@. Return these buffers as a list of ByteStrings |
---|
112 | createByteStringArray :: Int -- ^ the number of buffers requested |
---|
113 | -> Int -- ^ the size of each buffer |
---|
114 | -> ((Ptr (Ptr Word8)) -> IO ()) |
---|
115 | -> IO [B.ByteString] |
---|
116 | createByteStringArray n size f = do |
---|
117 | allocaBytes (n * sizeOf (undefined :: Ptr Word8)) (\array -> do |
---|
118 | allocaBytes (n * size) (\ptr -> do |
---|
119 | mapM_ (\i -> poke (advancePtr array i) (advancePtr ptr (size * i))) [0..(n - 1)] |
---|
120 | f array |
---|
121 | mapM (\i -> B.packCStringLen (castPtr $ advancePtr ptr (i * size), size)) [0..(n - 1)])) |
---|
122 | |
---|
123 | -- | Generate the secondary blocks from a list of the primary blocks. The |
---|
124 | -- primary blocks must be in order and all of the same size. There must be |
---|
125 | -- @k@ primary blocks. |
---|
126 | encode :: FECParams |
---|
127 | -> [B.ByteString] -- ^ a list of @k@ input blocks |
---|
128 | -> [B.ByteString] -- ^ (n - k) output blocks |
---|
129 | encode (FECParams params k n) inblocks |
---|
130 | | length inblocks /= k = error "Wrong number of blocks to FEC encode" |
---|
131 | | not (allByteStringsSameLength inblocks) = error "Not all inputs to FEC encode are the same length" |
---|
132 | | otherwise = unsafePerformIO (do |
---|
133 | let sz = B.length $ head inblocks |
---|
134 | withForeignPtr params (\cfec -> do |
---|
135 | byteStringsToArray inblocks (\src -> do |
---|
136 | createByteStringArray (n - k) sz (\fecs -> do |
---|
137 | uintCArray [k..(n - 1)] (\block_nums -> do |
---|
138 | _encode cfec src fecs block_nums (fromIntegral (n - k)) $ fromIntegral sz))))) |
---|
139 | |
---|
140 | -- | A sort function for tagged assoc lists |
---|
141 | sortTagged :: [(Int, a)] -> [(Int, a)] |
---|
142 | sortTagged = sortBy (\a b -> compare (fst a) (fst b)) |
---|
143 | |
---|
144 | -- | Reorder the given list so that elements with tag numbers < the first |
---|
145 | -- argument have an index equal to their tag number (if possible) |
---|
146 | reorderPrimaryBlocks :: Int -> [(Int, a)] -> [(Int, a)] |
---|
147 | reorderPrimaryBlocks n blocks = inner (sortTagged pBlocks) sBlocks [] where |
---|
148 | (pBlocks, sBlocks) = partition (\(tag, _) -> tag < n) blocks |
---|
149 | inner [] sBlocks acc = acc ++ sBlocks |
---|
150 | inner pBlocks [] acc = acc ++ pBlocks |
---|
151 | inner pBlocks@((tag, a) : ps) sBlocks@(s : ss) acc = |
---|
152 | if length acc == tag |
---|
153 | then inner ps sBlocks (acc ++ [(tag, a)]) |
---|
154 | else inner pBlocks ss (acc ++ [s]) |
---|
155 | |
---|
156 | -- | Recover the primary blocks from a list of @k@ blocks. Each block must be |
---|
157 | -- tagged with its number (see the module comments about block numbering) |
---|
158 | decode :: FECParams |
---|
159 | -> [(Int, B.ByteString)] -- ^ a list of @k@ blocks and their index |
---|
160 | -> [B.ByteString] -- ^ a list the @k@ primary blocks |
---|
161 | decode (FECParams params k n) inblocks |
---|
162 | | length (nub $ map fst inblocks) /= length (inblocks) = error "Duplicate input blocks in FEC decode" |
---|
163 | | any (\f -> f < 0 || f >= n) $ map fst inblocks = error "Invalid block numbers in FEC decode" |
---|
164 | | length inblocks /= k = error "Wrong number of blocks to FEC decode" |
---|
165 | | not (allByteStringsSameLength $ map snd inblocks) = error "Not all inputs to FEC decode are same length" |
---|
166 | | otherwise = unsafePerformIO (do |
---|
167 | let sz = B.length $ snd $ head inblocks |
---|
168 | inblocks' = reorderPrimaryBlocks k inblocks |
---|
169 | presentBlocks = map fst inblocks' |
---|
170 | withForeignPtr params (\cfec -> do |
---|
171 | byteStringsToArray (map snd inblocks') (\src -> do |
---|
172 | b <- createByteStringArray (n - k) sz (\out -> do |
---|
173 | uintCArray presentBlocks (\block_nums -> do |
---|
174 | _decode cfec src out block_nums $ fromIntegral sz)) |
---|
175 | let blocks = [0..(n - 1)] \\ presentBlocks |
---|
176 | tagged = zip blocks b |
---|
177 | allBlocks = sortTagged $ tagged ++ inblocks' |
---|
178 | return $ take k $ map snd allBlocks))) |
---|
179 | |
---|
180 | -- | Break a ByteString into @n@ parts, equal in length to the original, such |
---|
181 | -- that all @n@ are required to reconstruct the original, but having less |
---|
182 | -- than @n@ parts reveals no information about the orginal. |
---|
183 | -- |
---|
184 | -- This code works in IO monad because it needs a source of random bytes, |
---|
185 | -- which it gets from /dev/urandom. If this file doesn't exist an |
---|
186 | -- exception results |
---|
187 | -- |
---|
188 | -- Not terribly fast - probably best to do it with short inputs (e.g. an |
---|
189 | -- encryption key) |
---|
190 | secureDivide :: Int -- ^ the number of parts requested |
---|
191 | -> B.ByteString -- ^ the data to be split |
---|
192 | -> IO [B.ByteString] |
---|
193 | secureDivide n input |
---|
194 | | n < 0 = error "secureDivide called with negative number of parts" |
---|
195 | | otherwise = withFile "/dev/urandom" ReadMode (\handle -> do |
---|
196 | let inner 1 bs = return [bs] |
---|
197 | inner n bs = do |
---|
198 | mask <- B.hGet handle (B.length bs) |
---|
199 | let masked = B.pack $ B.zipWith xor bs mask |
---|
200 | rest <- inner (n - 1) masked |
---|
201 | return (mask : rest) |
---|
202 | inner n input) |
---|
203 | |
---|
204 | -- | Reverse the operation of secureDivide. The order of the inputs doesn't |
---|
205 | -- matter, but they must all be the same length |
---|
206 | secureCombine :: [B.ByteString] -> B.ByteString |
---|
207 | secureCombine [] = error "Passed empty list of inputs to secureCombine" |
---|
208 | secureCombine [a] = a |
---|
209 | secureCombine [a, b] = B.pack $ B.zipWith xor a b |
---|
210 | secureCombine (a : rest) = B.pack $ B.zipWith xor a $ secureCombine rest |
---|
211 | |
---|
212 | -- | A utility function which takes an arbitary input and FEC encodes it into a |
---|
213 | -- number of blocks. The order the resulting blocks doesn't matter so long |
---|
214 | -- as you have enough to present to @deFEC@. |
---|
215 | enFEC :: Int -- ^ the number of blocks required to reconstruct |
---|
216 | -> Int -- ^ the total number of blocks |
---|
217 | -> B.ByteString -- ^ the data to divide |
---|
218 | -> [B.ByteString] -- ^ the resulting blocks |
---|
219 | enFEC k n input = taggedPrimaryBlocks ++ taggedSecondaryBlocks where |
---|
220 | taggedPrimaryBlocks = map (uncurry B.cons) $ zip [0..] primaryBlocks |
---|
221 | taggedSecondaryBlocks = map (uncurry B.cons) $ zip [(fromIntegral k)..] secondaryBlocks |
---|
222 | remainder = B.length input `mod` k |
---|
223 | paddingLength = if remainder >= 1 then (k - remainder) else k |
---|
224 | paddingBytes = (B.replicate (paddingLength - 1) 0) `B.append` (B.singleton $ fromIntegral paddingLength) |
---|
225 | divide a bs |
---|
226 | | B.null bs = [] |
---|
227 | | otherwise = (B.take a bs) : (divide a $ B.drop a bs) |
---|
228 | input' = input `B.append` paddingBytes |
---|
229 | blockSize = B.length input' `div` k |
---|
230 | primaryBlocks = divide blockSize input' |
---|
231 | secondaryBlocks = encode params primaryBlocks |
---|
232 | params = fec k n |
---|
233 | |
---|
234 | -- | Reverses the operation of @enFEC@. |
---|
235 | deFEC :: Int -- ^ the number of blocks required (matches call to @enFEC@) |
---|
236 | -> Int -- ^ the total number of blocks (matches call to @enFEC@) |
---|
237 | -> [B.ByteString] -- ^ a list of k, or more, blocks from @enFEC@ |
---|
238 | -> B.ByteString |
---|
239 | deFEC k n inputs |
---|
240 | | length inputs < k = error "Too few inputs to deFEC" |
---|
241 | | otherwise = B.take (B.length fecOutput - paddingLength) fecOutput where |
---|
242 | paddingLength = fromIntegral $ B.last fecOutput |
---|
243 | inputs' = take k inputs |
---|
244 | taggedInputs = map (\bs -> (fromIntegral $ B.head bs, B.tail bs)) inputs' |
---|
245 | fecOutput = B.concat $ decode params taggedInputs |
---|
246 | params = fec k n |
---|