1 {-# LANGUAGE ForeignFunctionInterface, EmptyDataDecls #-}
4 -- Copyright: Adam Langley
5 -- License: GPLv2+|TGPPLv1+ (see README.rst for details)
7 -- Stability: experimental
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.
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.
23 -- * Utility functions
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)
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)
46 data FECParams = FECParams (ForeignPtr CFEC) Int Int
48 instance Show FECParams where
49 show (FECParams _ k n) = "FEC (" ++ show k ++ ", " ++ show n ++ ")"
51 foreign import ccall unsafe "fec_new" _new :: CUInt -- ^ k
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
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
69 -- | Return true if the given @k@ and @n@ values are valid
70 isValidConfig :: Int -> Int -> Bool
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
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)
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
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
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)
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
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 ())
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)]
121 mapM (\i -> B.packCStringLen (castPtr $ advancePtr ptr (i * size), size)) [0..(n - 1)]))
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.
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)))))
140 -- | A sort function for tagged assoc lists
141 sortTagged :: [(Int, a)] -> [(Int, a)]
142 sortTagged = sortBy (\a b -> compare (fst a) (fst b))
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 =
153 then inner ps sBlocks (acc ++ [(tag, a)])
154 else inner pBlocks ss (acc ++ [s])
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)
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)))
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.
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
188 -- Not terribly fast - probably best to do it with short inputs (e.g. an
190 secureDivide :: Int -- ^ the number of parts requested
191 -> B.ByteString -- ^ the data to be split
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]
198 mask <- B.hGet handle (B.length bs)
199 let masked = B.pack $ B.zipWith xor bs mask
200 rest <- inner (n - 1) masked
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
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)
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
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@
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