zfec/haskell/Codec/FEC.hs

   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