module Codec.FEC (
FECParams
, fec
, encode
, decode
, 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
-> CUInt
-> 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)
-> Ptr (Ptr Word8)
-> Ptr CUInt
-> CSize
-> CSize
-> IO ()
foreign import ccall unsafe "fec_decode" _decode :: Ptr CFEC
-> Ptr (Ptr Word8)
-> Ptr (Ptr Word8)
-> Ptr CUInt
-> CSize
-> IO ()
isValidConfig :: Int -> Int -> Bool
isValidConfig k n
| k >= n = False
| k < 1 = False
| n < 1 = False
| n > 255 = False
| otherwise = True
fec :: Int
-> Int
-> 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)
uintCArray :: [Int] -> ((Ptr CUInt) -> IO a) -> IO a
uintCArray xs f = withArray (map fromIntegral xs) f
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)
allByteStringsSameLength :: [B.ByteString] -> Bool
allByteStringsSameLength [] = True
allByteStringsSameLength (bs : bss) = all ((==) (B.length bs)) $ map B.length bss
createByteStringArray :: Int
-> Int
-> ((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)]))
encode :: FECParams
-> [B.ByteString]
-> [B.ByteString]
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)))))
sortTagged :: [(Int, a)] -> [(Int, a)]
sortTagged = sortBy (\a b -> compare (fst a) (fst b))
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])
decode :: FECParams
-> [(Int, B.ByteString)]
-> [B.ByteString]
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)))
secureDivide :: Int
-> B.ByteString
-> 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)
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
enFEC :: Int
-> Int
-> B.ByteString
-> [B.ByteString]
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
deFEC :: Int
-> Int
-> [B.ByteString]
-> 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