+++ /dev/null
-{-# LANGUAGE ForeignFunctionInterface, EmptyDataDecls #-}
--- |
--- Module: Codec.FEC
--- Copyright: Adam Langley
--- License: GPLv2+|TGPPLv1+ (see README.rst for details)
---
--- Stability: experimental
---
--- The module provides k of n encoding - a way to generate (n - k) secondary
--- blocks of data from k primary blocks such that any k blocks (primary or
--- secondary) are sufficient to regenerate all blocks.
---
--- All blocks must be the same length and you need to keep track of which
--- blocks you have in order to tell decode. By convention, the blocks are
--- numbered 0..(n - 1) and blocks numbered < k are the primary blocks.
-
-module Codec.FEC (
- FECParams
- , fec
- , encode
- , decode
-
- -- * Utility functions
- , secureDivide
- , secureCombine
- , enFEC
- , deFEC
- ) where
-
-import qualified Data.ByteString as B
-import qualified Data.ByteString.Unsafe as BU
-import qualified Data.ByteString.Internal as BI
-import Data.Word (Word8)
-import Data.Bits (xor)
-import Data.List (sortBy, partition, (\\), nub)
-import Foreign.Ptr
-import Foreign.Storable (sizeOf, poke)
-import Foreign.ForeignPtr
-import Foreign.C.Types
-import Foreign.Marshal.Alloc
-import Foreign.Marshal.Array (withArray, advancePtr)
-import System.IO (withFile, IOMode(..))
-import System.IO.Unsafe (unsafePerformIO)
-
-data CFEC
-data FECParams = FECParams (ForeignPtr CFEC) Int Int
-
-instance Show FECParams where
- show (FECParams _ k n) = "FEC (" ++ show k ++ ", " ++ show n ++ ")"
-
-foreign import ccall unsafe "fec_new" _new :: CUInt -- ^ k
- -> CUInt -- ^ n
- -> IO (Ptr CFEC)
-foreign import ccall unsafe "&fec_free" _free :: FunPtr (Ptr CFEC -> IO ())
-foreign import ccall unsafe "fec_encode" _encode :: Ptr CFEC
- -> Ptr (Ptr Word8) -- ^ primary blocks
- -> Ptr (Ptr Word8) -- ^ (output) secondary blocks
- -> Ptr CUInt -- ^ array of secondary block ids
- -> CSize -- ^ length of previous
- -> CSize -- ^ block length
- -> IO ()
-foreign import ccall unsafe "fec_decode" _decode :: Ptr CFEC
- -> Ptr (Ptr Word8) -- ^ input blocks
- -> Ptr (Ptr Word8) -- ^ output blocks
- -> Ptr CUInt -- ^ array of input indexes
- -> CSize -- ^ block length
- -> IO ()
-
--- | Return true if the given @k@ and @n@ values are valid
-isValidConfig :: Int -> Int -> Bool
-isValidConfig k n
- | k >= n = False
- | k < 1 = False
- | n < 1 = False
- | n > 255 = False
- | otherwise = True
-
--- | Return a FEC with the given parameters.
-fec :: Int -- ^ the number of primary blocks
- -> Int -- ^ the total number blocks, must be < 256
- -> FECParams
-fec k n =
- if not (isValidConfig k n)
- then error $ "Invalid FEC parameters: " ++ show k ++ " " ++ show n
- else unsafePerformIO (do
- cfec <- _new (fromIntegral k) (fromIntegral n)
- params <- newForeignPtr _free cfec
- return $ FECParams params k n)
-
--- | Create a C array of unsigned from an input array
-uintCArray :: [Int] -> ((Ptr CUInt) -> IO a) -> IO a
-uintCArray xs f = withArray (map fromIntegral xs) f
-
--- | Convert a list of ByteStrings to an array of pointers to their data
-byteStringsToArray :: [B.ByteString] -> ((Ptr (Ptr Word8)) -> IO a) -> IO a
-byteStringsToArray inputs f = do
- let l = length inputs
- allocaBytes (l * sizeOf (undefined :: Ptr Word8)) (\array -> do
- let inner _ [] = f array
- inner array' (bs : bss) = BU.unsafeUseAsCString bs (\ptr -> do
- poke array' $ castPtr ptr
- inner (advancePtr array' 1) bss)
- inner array inputs)
-
--- | Return True iff all the given ByteStrings are the same length
-allByteStringsSameLength :: [B.ByteString] -> Bool
-allByteStringsSameLength [] = True
-allByteStringsSameLength (bs : bss) = all ((==) (B.length bs)) $ map B.length bss
-
--- | Run the given function with a pointer to an array of @n@ pointers to
--- buffers of size @size@. Return these buffers as a list of ByteStrings
-createByteStringArray :: Int -- ^ the number of buffers requested
- -> Int -- ^ the size of each buffer
- -> ((Ptr (Ptr Word8)) -> IO ())
- -> IO [B.ByteString]
-createByteStringArray n size f = do
- allocaBytes (n * sizeOf (undefined :: Ptr Word8)) (\array -> do
- allocaBytes (n * size) (\ptr -> do
- mapM_ (\i -> poke (advancePtr array i) (advancePtr ptr (size * i))) [0..(n - 1)]
- f array
- mapM (\i -> B.packCStringLen (castPtr $ advancePtr ptr (i * size), size)) [0..(n - 1)]))
-
--- | Generate the secondary blocks from a list of the primary blocks. The
--- primary blocks must be in order and all of the same size. There must be
--- @k@ primary blocks.
-encode :: FECParams
- -> [B.ByteString] -- ^ a list of @k@ input blocks
- -> [B.ByteString] -- ^ (n - k) output blocks
-encode (FECParams params k n) inblocks
- | length inblocks /= k = error "Wrong number of blocks to FEC encode"
- | not (allByteStringsSameLength inblocks) = error "Not all inputs to FEC encode are the same length"
- | otherwise = unsafePerformIO (do
- let sz = B.length $ head inblocks
- withForeignPtr params (\cfec -> do
- byteStringsToArray inblocks (\src -> do
- createByteStringArray (n - k) sz (\fecs -> do
- uintCArray [k..(n - 1)] (\block_nums -> do
- _encode cfec src fecs block_nums (fromIntegral (n - k)) $ fromIntegral sz)))))
-
--- | A sort function for tagged assoc lists
-sortTagged :: [(Int, a)] -> [(Int, a)]
-sortTagged = sortBy (\a b -> compare (fst a) (fst b))
-
--- | Reorder the given list so that elements with tag numbers < the first
--- argument have an index equal to their tag number (if possible)
-reorderPrimaryBlocks :: Int -> [(Int, a)] -> [(Int, a)]
-reorderPrimaryBlocks n blocks = inner (sortTagged pBlocks) sBlocks [] where
- (pBlocks, sBlocks) = partition (\(tag, _) -> tag < n) blocks
- inner [] sBlocks acc = acc ++ sBlocks
- inner pBlocks [] acc = acc ++ pBlocks
- inner pBlocks@((tag, a) : ps) sBlocks@(s : ss) acc =
- if length acc == tag
- then inner ps sBlocks (acc ++ [(tag, a)])
- else inner pBlocks ss (acc ++ [s])
-
--- | Recover the primary blocks from a list of @k@ blocks. Each block must be
--- tagged with its number (see the module comments about block numbering)
-decode :: FECParams
- -> [(Int, B.ByteString)] -- ^ a list of @k@ blocks and their index
- -> [B.ByteString] -- ^ a list the @k@ primary blocks
-decode (FECParams params k n) inblocks
- | length (nub $ map fst inblocks) /= length (inblocks) = error "Duplicate input blocks in FEC decode"
- | any (\f -> f < 0 || f >= n) $ map fst inblocks = error "Invalid block numbers in FEC decode"
- | length inblocks /= k = error "Wrong number of blocks to FEC decode"
- | not (allByteStringsSameLength $ map snd inblocks) = error "Not all inputs to FEC decode are same length"
- | otherwise = unsafePerformIO (do
- let sz = B.length $ snd $ head inblocks
- inblocks' = reorderPrimaryBlocks k inblocks
- presentBlocks = map fst inblocks'
- withForeignPtr params (\cfec -> do
- byteStringsToArray (map snd inblocks') (\src -> do
- b <- createByteStringArray (n - k) sz (\out -> do
- uintCArray presentBlocks (\block_nums -> do
- _decode cfec src out block_nums $ fromIntegral sz))
- let blocks = [0..(n - 1)] \\ presentBlocks
- tagged = zip blocks b
- allBlocks = sortTagged $ tagged ++ inblocks'
- return $ take k $ map snd allBlocks)))
-
--- | Break a ByteString into @n@ parts, equal in length to the original, such
--- that all @n@ are required to reconstruct the original, but having less
--- than @n@ parts reveals no information about the orginal.
---
--- This code works in IO monad because it needs a source of random bytes,
--- which it gets from /dev/urandom. If this file doesn't exist an
--- exception results
---
--- Not terribly fast - probably best to do it with short inputs (e.g. an
--- encryption key)
-secureDivide :: Int -- ^ the number of parts requested
- -> B.ByteString -- ^ the data to be split
- -> IO [B.ByteString]
-secureDivide n input
- | n < 0 = error "secureDivide called with negative number of parts"
- | otherwise = withFile "/dev/urandom" ReadMode (\handle -> do
- let inner 1 bs = return [bs]
- inner n bs = do
- mask <- B.hGet handle (B.length bs)
- let masked = B.pack $ B.zipWith xor bs mask
- rest <- inner (n - 1) masked
- return (mask : rest)
- inner n input)
-
--- | Reverse the operation of secureDivide. The order of the inputs doesn't
--- matter, but they must all be the same length
-secureCombine :: [B.ByteString] -> B.ByteString
-secureCombine [] = error "Passed empty list of inputs to secureCombine"
-secureCombine [a] = a
-secureCombine [a, b] = B.pack $ B.zipWith xor a b
-secureCombine (a : rest) = B.pack $ B.zipWith xor a $ secureCombine rest
-
--- | A utility function which takes an arbitary input and FEC encodes it into a
--- number of blocks. The order the resulting blocks doesn't matter so long
--- as you have enough to present to @deFEC@.
-enFEC :: Int -- ^ the number of blocks required to reconstruct
- -> Int -- ^ the total number of blocks
- -> B.ByteString -- ^ the data to divide
- -> [B.ByteString] -- ^ the resulting blocks
-enFEC k n input = taggedPrimaryBlocks ++ taggedSecondaryBlocks where
- taggedPrimaryBlocks = map (uncurry B.cons) $ zip [0..] primaryBlocks
- taggedSecondaryBlocks = map (uncurry B.cons) $ zip [(fromIntegral k)..] secondaryBlocks
- remainder = B.length input `mod` k
- paddingLength = if remainder >= 1 then (k - remainder) else k
- paddingBytes = (B.replicate (paddingLength - 1) 0) `B.append` (B.singleton $ fromIntegral paddingLength)
- divide a bs
- | B.null bs = []
- | otherwise = (B.take a bs) : (divide a $ B.drop a bs)
- input' = input `B.append` paddingBytes
- blockSize = B.length input' `div` k
- primaryBlocks = divide blockSize input'
- secondaryBlocks = encode params primaryBlocks
- params = fec k n
-
--- | Reverses the operation of @enFEC@.
-deFEC :: Int -- ^ the number of blocks required (matches call to @enFEC@)
- -> Int -- ^ the total number of blocks (matches call to @enFEC@)
- -> [B.ByteString] -- ^ a list of k, or more, blocks from @enFEC@
- -> B.ByteString
-deFEC k n inputs
- | length inputs < k = error "Too few inputs to deFEC"
- | otherwise = B.take (B.length fecOutput - paddingLength) fecOutput where
- paddingLength = fromIntegral $ B.last fecOutput
- inputs' = take k inputs
- taggedInputs = map (\bs -> (fromIntegral $ B.head bs, B.tail bs)) inputs'
- fecOutput = B.concat $ decode params taggedInputs
- params = fec k n
projects / tahoe-lafs / zfec.git / blobdiff