{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE NamedFieldPuns #-}

{- |
 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 (paramK, paramN),
    initialize,
    fec,
    encode,
    decode,

    -- * Utility functions
    secureDivide,
    secureCombine,
    enFEC,
    deFEC,
) where

import Control.Concurrent.Extra (Lock, newLock, withLock)
import Control.DeepSeq (NFData (rnf))
import Control.Exception (Exception, throwIO)
import Data.Bits (xor)
import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as BU
import Data.List (nub, partition, sortBy, (\\))
import Data.Word (Word8)
import Foreign.C.Types (CSize (..), CUInt (..))
import Foreign.ForeignPtr (
    ForeignPtr,
    newForeignPtr,
    withForeignPtr,
 )
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Marshal.Array (advancePtr, withArray)
import Foreign.Ptr (FunPtr, Ptr, castPtr, nullPtr)
import Foreign.Storable (poke, sizeOf)
import GHC.Generics (Generic)
import System.IO (IOMode (..), withFile)
import System.IO.Unsafe (unsafePerformIO)

data CFEC
data FECParams = FECParams
    { FECParams -> ForeignPtr CFEC
_cfec :: !(ForeignPtr CFEC)
    , FECParams -> Int
paramK :: Int
    , FECParams -> Int
paramN :: Int
    }
    deriving ((forall x. FECParams -> Rep FECParams x)
-> (forall x. Rep FECParams x -> FECParams) -> Generic FECParams
forall x. Rep FECParams x -> FECParams
forall x. FECParams -> Rep FECParams x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep FECParams x -> FECParams
$cfrom :: forall x. FECParams -> Rep FECParams x
Generic)

-- Provide an NFData instance so it's possible to use a FECParams in a
-- Criterion benchmark.
instance NFData FECParams where
    rnf :: FECParams -> ()
rnf FECParams{ForeignPtr CFEC
_cfec :: ForeignPtr CFEC
_cfec :: FECParams -> ForeignPtr CFEC
_cfec, Int
paramK :: Int
paramK :: FECParams -> Int
paramK, Int
paramN :: Int
paramN :: FECParams -> Int
paramN} =
        -- ForeignPtr has no NFData instance and I don't know how to implement
        -- one for it so we punt on it here.  We do make it strict in the
        -- record definition which at least shallowly evaluates the
        -- ForeignPtr which is ... part of the job?
        Int -> ()
forall a. NFData a => a -> ()
rnf Int
paramK () -> () -> ()
`seq` Int -> ()
forall a. NFData a => a -> ()
rnf Int
paramN

instance Show FECParams where
    show :: FECParams -> String
show (FECParams ForeignPtr CFEC
_ Int
k Int
n) = String
"FEC (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

foreign import ccall unsafe "fec_init"
    _init :: IO ()
foreign import ccall unsafe "fec_new"
    _new ::
        -- | k
        CUInt ->
        -- | n
        CUInt ->
        IO (Ptr CFEC)
foreign import ccall unsafe "&fec_free" _free :: FunPtr (Ptr CFEC -> IO ())
foreign import ccall unsafe "fec_encode"
    _encode ::
        Ptr CFEC ->
        -- | primary blocks
        Ptr (Ptr Word8) ->
        -- | (output) secondary blocks
        Ptr (Ptr Word8) ->
        -- | array of secondary block ids
        Ptr CUInt ->
        -- | length of previous
        CSize ->
        -- | block length
        CSize ->
        IO ()
foreign import ccall unsafe "fec_decode"
    _decode ::
        Ptr CFEC ->
        -- | input blocks
        Ptr (Ptr Word8) ->
        -- | output blocks
        Ptr (Ptr Word8) ->
        -- | array of input indexes
        Ptr CUInt ->
        -- | block length
        CSize ->
        IO ()

-- | Return true if the given @k@ and @n@ values are valid
isValidConfig :: Int -> Int -> Bool
isValidConfig :: Int -> Int -> Bool
isValidConfig Int
k Int
n
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n = Bool
False
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = Bool
False
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = Bool
False
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
255 = Bool
False
    | Bool
otherwise = Bool
True

{- | The underlying library signaled that it has not been properly initialized
 yet.  Use @initialize@ to initialize it.
-}
data Uninitialized = Uninitialized deriving (Eq Uninitialized
Eq Uninitialized
-> (Uninitialized -> Uninitialized -> Ordering)
-> (Uninitialized -> Uninitialized -> Bool)
-> (Uninitialized -> Uninitialized -> Bool)
-> (Uninitialized -> Uninitialized -> Bool)
-> (Uninitialized -> Uninitialized -> Bool)
-> (Uninitialized -> Uninitialized -> Uninitialized)
-> (Uninitialized -> Uninitialized -> Uninitialized)
-> Ord Uninitialized
Uninitialized -> Uninitialized -> Bool
Uninitialized -> Uninitialized -> Ordering
Uninitialized -> Uninitialized -> Uninitialized
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Uninitialized -> Uninitialized -> Uninitialized
$cmin :: Uninitialized -> Uninitialized -> Uninitialized
max :: Uninitialized -> Uninitialized -> Uninitialized
$cmax :: Uninitialized -> Uninitialized -> Uninitialized
>= :: Uninitialized -> Uninitialized -> Bool
$c>= :: Uninitialized -> Uninitialized -> Bool
> :: Uninitialized -> Uninitialized -> Bool
$c> :: Uninitialized -> Uninitialized -> Bool
<= :: Uninitialized -> Uninitialized -> Bool
$c<= :: Uninitialized -> Uninitialized -> Bool
< :: Uninitialized -> Uninitialized -> Bool
$c< :: Uninitialized -> Uninitialized -> Bool
compare :: Uninitialized -> Uninitialized -> Ordering
$ccompare :: Uninitialized -> Uninitialized -> Ordering
$cp1Ord :: Eq Uninitialized
Ord, Uninitialized -> Uninitialized -> Bool
(Uninitialized -> Uninitialized -> Bool)
-> (Uninitialized -> Uninitialized -> Bool) -> Eq Uninitialized
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Uninitialized -> Uninitialized -> Bool
$c/= :: Uninitialized -> Uninitialized -> Bool
== :: Uninitialized -> Uninitialized -> Bool
$c== :: Uninitialized -> Uninitialized -> Bool
Eq, Int -> Uninitialized -> ShowS
[Uninitialized] -> ShowS
Uninitialized -> String
(Int -> Uninitialized -> ShowS)
-> (Uninitialized -> String)
-> ([Uninitialized] -> ShowS)
-> Show Uninitialized
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Uninitialized] -> ShowS
$cshowList :: [Uninitialized] -> ShowS
show :: Uninitialized -> String
$cshow :: Uninitialized -> String
showsPrec :: Int -> Uninitialized -> ShowS
$cshowsPrec :: Int -> Uninitialized -> ShowS
Show)

instance Exception Uninitialized

-- A lock to ensure at most one thread attempts to initialize the underlying
-- library at a time.  Multiple initializations are harmless but concurrent
-- initializations are disallowed.
_initializationLock :: Lock
{-# NOINLINE _initializationLock #-}
_initializationLock :: Lock
_initializationLock = IO Lock -> Lock
forall a. IO a -> a
unsafePerformIO IO Lock
newLock

-- | Initialize the library.  This must be done before other APIs can succeed.
initialize :: IO ()
initialize :: IO ()
initialize = Lock -> IO () -> IO ()
forall a. Lock -> IO a -> IO a
withLock Lock
_initializationLock IO ()
_init

-- | Return a FEC with the given parameters.
fec ::
    -- | the number of primary blocks
    Int ->
    -- | the total number blocks, must be < 256
    Int ->
    FECParams
fec :: Int -> Int -> FECParams
fec Int
k Int
n =
    if Bool -> Bool
not (Int -> Int -> Bool
isValidConfig Int
k Int
n)
        then String -> FECParams
forall a. HasCallStack => String -> a
error (String -> FECParams) -> String -> FECParams
forall a b. (a -> b) -> a -> b
$ String
"Invalid FEC parameters: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
        else
            IO FECParams -> FECParams
forall a. IO a -> a
unsafePerformIO
                ( do
                    Ptr CFEC
cfec' <- CUInt -> CUInt -> IO (Ptr CFEC)
_new (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
                    -- new will return null if the library hasn't been
                    -- initialized.
                    if Ptr CFEC
cfec' Ptr CFEC -> Ptr CFEC -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr CFEC
forall a. Ptr a
nullPtr
                        then Uninitialized -> IO FECParams
forall e a. Exception e => e -> IO a
throwIO Uninitialized
Uninitialized
                        else do
                            ForeignPtr CFEC
params <- FinalizerPtr CFEC -> Ptr CFEC -> IO (ForeignPtr CFEC)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr CFEC
_free Ptr CFEC
cfec'
                            FECParams -> IO FECParams
forall (m :: * -> *) a. Monad m => a -> m a
return (FECParams -> IO FECParams) -> FECParams -> IO FECParams
forall a b. (a -> b) -> a -> b
$ ForeignPtr CFEC -> Int -> Int -> FECParams
FECParams ForeignPtr CFEC
params Int
k Int
n
                )

-- | Create a C array of unsigned from an input array
uintCArray :: [Int] -> (Ptr CUInt -> IO a) -> IO a
uintCArray :: [Int] -> (Ptr CUInt -> IO a) -> IO a
uintCArray = [CUInt] -> (Ptr CUInt -> IO a) -> IO a
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray ([CUInt] -> (Ptr CUInt -> IO a) -> IO a)
-> ([Int] -> [CUInt]) -> [Int] -> (Ptr CUInt -> IO a) -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> CUInt) -> [Int] -> [CUInt]
forall a b. (a -> b) -> [a] -> [b]
map Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | Convert a list of ByteStrings to an array of pointers to their data
byteStringsToArray :: [B.ByteString] -> (Ptr (Ptr Word8) -> IO a) -> IO a
byteStringsToArray :: [ByteString] -> (Ptr (Ptr Word8) -> IO a) -> IO a
byteStringsToArray [ByteString]
inputs Ptr (Ptr Word8) -> IO a
f = do
    let l :: Int
l = [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
inputs
    Int -> (Ptr (Ptr Word8) -> IO a) -> IO a
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes
        (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
* Ptr Word8 -> Int
forall a. Storable a => a -> Int
sizeOf (Ptr Word8
forall a. HasCallStack => a
undefined :: Ptr Word8))
        ( \Ptr (Ptr Word8)
array -> do
            let inner :: Ptr (Ptr b) -> [ByteString] -> IO a
inner Ptr (Ptr b)
_ [] = Ptr (Ptr Word8) -> IO a
f Ptr (Ptr Word8)
array
                inner Ptr (Ptr b)
array' (ByteString
bs : [ByteString]
bss) =
                    ByteString -> (CString -> IO a) -> IO a
forall a. ByteString -> (CString -> IO a) -> IO a
BU.unsafeUseAsCString
                        ByteString
bs
                        ( \CString
ptr -> do
                            Ptr (Ptr b) -> Ptr b -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (Ptr b)
array' (Ptr b -> IO ()) -> Ptr b -> IO ()
forall a b. (a -> b) -> a -> b
$ CString -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr CString
ptr
                            Ptr (Ptr b) -> [ByteString] -> IO a
inner (Ptr (Ptr b) -> Int -> Ptr (Ptr b)
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr (Ptr b)
array' Int
1) [ByteString]
bss
                        )
            Ptr (Ptr Word8) -> [ByteString] -> IO a
forall b. Ptr (Ptr b) -> [ByteString] -> IO a
inner Ptr (Ptr Word8)
array [ByteString]
inputs
        )

-- | Return True iff all the given ByteStrings are the same length
allByteStringsSameLength :: [B.ByteString] -> Bool
allByteStringsSameLength :: [ByteString] -> Bool
allByteStringsSameLength [] = Bool
True
allByteStringsSameLength (ByteString
bs : [ByteString]
bss) = (ByteString -> Bool) -> [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (ByteString -> Int
B.length ByteString
bs) (Int -> Bool) -> (ByteString -> Int) -> ByteString -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Int
B.length) [ByteString]
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 ::
    -- | the number of buffers requested
    Int ->
    -- | the size of each buffer
    Int ->
    (Ptr (Ptr Word8) -> IO ()) ->
    IO [B.ByteString]
createByteStringArray :: Int -> Int -> (Ptr (Ptr Word8) -> IO ()) -> IO [ByteString]
createByteStringArray Int
n Int
size Ptr (Ptr Word8) -> IO ()
f = do
    Int -> (Ptr (Ptr Word8) -> IO [ByteString]) -> IO [ByteString]
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes
        (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Ptr Word8 -> Int
forall a. Storable a => a -> Int
sizeOf (Ptr Word8
forall a. HasCallStack => a
undefined :: Ptr Word8))
        ( \Ptr (Ptr Word8)
array -> do
            Int -> (Ptr Word8 -> IO [ByteString]) -> IO [ByteString]
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes
                (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)
                ( \Ptr Word8
ptr -> do
                    (Int -> IO ()) -> [Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Int
i -> Ptr (Ptr Word8) -> Ptr Word8 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr (Ptr Word8) -> Int -> Ptr (Ptr Word8)
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr (Ptr Word8)
array Int
i) (Ptr Word8 -> Int -> Ptr Word8
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr Word8
ptr (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i))) [Int
0 .. (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
                    Ptr (Ptr Word8) -> IO ()
f Ptr (Ptr Word8)
array
                    (Int -> IO ByteString) -> [Int] -> IO [ByteString]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Int
i -> CStringLen -> IO ByteString
B.packCStringLen (Ptr Word8 -> CString
forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8 -> CString) -> Ptr Word8 -> CString
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> Ptr Word8
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr Word8
ptr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size), Int
size)) [Int
0 .. (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 ->
    -- | a list of @k@ input blocks
    [B.ByteString] ->
    -- | (n - k) output blocks
    [B.ByteString]
encode :: FECParams -> [ByteString] -> [ByteString]
encode (FECParams ForeignPtr CFEC
params Int
k Int
n) [ByteString]
inblocks
    | [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
inblocks Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
k = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Wrong number of blocks to FEC encode"
    | Bool -> Bool
not ([ByteString] -> Bool
allByteStringsSameLength [ByteString]
inblocks) = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Not all inputs to FEC encode are the same length"
    | Bool
otherwise =
        IO [ByteString] -> [ByteString]
forall a. IO a -> a
unsafePerformIO
            ( do
                let sz :: Int
sz = ByteString -> Int
B.length (ByteString -> Int) -> ByteString -> Int
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
forall a. [a] -> a
head [ByteString]
inblocks
                ForeignPtr CFEC -> (Ptr CFEC -> IO [ByteString]) -> IO [ByteString]
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr
                    ForeignPtr CFEC
params
                    ( \Ptr CFEC
cfec' -> do
                        [ByteString]
-> (Ptr (Ptr Word8) -> IO [ByteString]) -> IO [ByteString]
forall a. [ByteString] -> (Ptr (Ptr Word8) -> IO a) -> IO a
byteStringsToArray
                            [ByteString]
inblocks
                            ( \Ptr (Ptr Word8)
src -> do
                                Int -> Int -> (Ptr (Ptr Word8) -> IO ()) -> IO [ByteString]
createByteStringArray
                                    (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)
                                    Int
sz
                                    ( \Ptr (Ptr Word8)
fecs -> do
                                        [Int] -> (Ptr CUInt -> IO ()) -> IO ()
forall a. [Int] -> (Ptr CUInt -> IO a) -> IO a
uintCArray
                                            [Int
k .. (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
                                            ( \Ptr CUInt
block_nums -> do
                                                Ptr CFEC
-> Ptr (Ptr Word8)
-> Ptr (Ptr Word8)
-> Ptr CUInt
-> CSize
-> CSize
-> IO ()
_encode Ptr CFEC
cfec' Ptr (Ptr Word8)
src Ptr (Ptr Word8)
fecs Ptr CUInt
block_nums (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)) (CSize -> IO ()) -> CSize -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz
                                            )
                                    )
                            )
                    )
            )

-- | A sort function for tagged assoc lists
sortTagged :: [(Int, a)] -> [(Int, a)]
sortTagged :: [(Int, a)] -> [(Int, a)]
sortTagged = ((Int, a) -> (Int, a) -> Ordering) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (\(Int, a)
a (Int, a)
b -> Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ((Int, a) -> Int
forall a b. (a, b) -> a
fst (Int, a)
a) ((Int, a) -> Int
forall a b. (a, b) -> a
fst (Int, a)
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 :: Int -> [(Int, a)] -> [(Int, a)]
reorderPrimaryBlocks Int
n [(Int, a)]
blocks = [(Int, a)] -> [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall b. [(Int, b)] -> [(Int, b)] -> [(Int, b)] -> [(Int, b)]
inner ([(Int, a)] -> [(Int, a)]
forall a. [(Int, a)] -> [(Int, a)]
sortTagged [(Int, a)]
pBlocks) [(Int, a)]
sBlocks []
  where
    ([(Int, a)]
pBlocks, [(Int, a)]
sBlocks) = ((Int, a) -> Bool) -> [(Int, a)] -> ([(Int, a)], [(Int, a)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (\(Int
tag, a
_) -> Int
tag Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) [(Int, a)]
blocks
    inner :: [(Int, b)] -> [(Int, b)] -> [(Int, b)] -> [(Int, b)]
inner [] [(Int, b)]
sBlocks' [(Int, b)]
acc = [(Int, b)]
acc [(Int, b)] -> [(Int, b)] -> [(Int, b)]
forall a. [a] -> [a] -> [a]
++ [(Int, b)]
sBlocks'
    inner [(Int, b)]
pBlocks' [] [(Int, b)]
acc = [(Int, b)]
acc [(Int, b)] -> [(Int, b)] -> [(Int, b)]
forall a. [a] -> [a] -> [a]
++ [(Int, b)]
pBlocks'
    inner pBlocks' :: [(Int, b)]
pBlocks'@((Int
tag, b
a) : [(Int, b)]
ps) sBlocks' :: [(Int, b)]
sBlocks'@((Int, b)
s : [(Int, b)]
ss) [(Int, b)]
acc =
        if [(Int, b)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, b)]
acc Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
tag
            then [(Int, b)] -> [(Int, b)] -> [(Int, b)] -> [(Int, b)]
inner [(Int, b)]
ps [(Int, b)]
sBlocks' ([(Int, b)]
acc [(Int, b)] -> [(Int, b)] -> [(Int, b)]
forall a. [a] -> [a] -> [a]
++ [(Int
tag, b
a)])
            else [(Int, b)] -> [(Int, b)] -> [(Int, b)] -> [(Int, b)]
inner [(Int, b)]
pBlocks' [(Int, b)]
ss ([(Int, b)]
acc [(Int, b)] -> [(Int, b)] -> [(Int, b)]
forall a. [a] -> [a] -> [a]
++ [(Int, b)
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 ->
    -- | a list of @k@ blocks and their index
    [(Int, B.ByteString)] ->
    -- | a list the @k@ primary blocks
    [B.ByteString]
decode :: FECParams -> [(Int, ByteString)] -> [ByteString]
decode (FECParams ForeignPtr CFEC
params Int
k Int
n) [(Int, ByteString)]
inblocks
    | [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ((Int, ByteString) -> Int) -> [(Int, ByteString)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst [(Int, ByteString)]
inblocks) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [(Int, ByteString)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, ByteString)]
inblocks = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Duplicate input blocks in FEC decode"
    | ((Int, ByteString) -> Bool) -> [(Int, ByteString)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((\Int
f -> Int
f Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
f Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n) (Int -> Bool)
-> ((Int, ByteString) -> Int) -> (Int, ByteString) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst) [(Int, ByteString)]
inblocks = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Invalid block numbers in FEC decode"
    | [(Int, ByteString)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, ByteString)]
inblocks Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
k = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Wrong number of blocks to FEC decode"
    | Bool -> Bool
not ([ByteString] -> Bool
allByteStringsSameLength ([ByteString] -> Bool) -> [ByteString] -> Bool
forall a b. (a -> b) -> a -> b
$ ((Int, ByteString) -> ByteString)
-> [(Int, ByteString)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd [(Int, ByteString)]
inblocks) = String -> [ByteString]
forall a. HasCallStack => String -> a
error String
"Not all inputs to FEC decode are same length"
    | Bool
otherwise =
        IO [ByteString] -> [ByteString]
forall a. IO a -> a
unsafePerformIO
            ( do
                let sz :: Int
sz = ByteString -> Int
B.length (ByteString -> Int) -> ByteString -> Int
forall a b. (a -> b) -> a -> b
$ (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd ((Int, ByteString) -> ByteString)
-> (Int, ByteString) -> ByteString
forall a b. (a -> b) -> a -> b
$ [(Int, ByteString)] -> (Int, ByteString)
forall a. [a] -> a
head [(Int, ByteString)]
inblocks
                    inblocks' :: [(Int, ByteString)]
inblocks' = Int -> [(Int, ByteString)] -> [(Int, ByteString)]
forall a. Int -> [(Int, a)] -> [(Int, a)]
reorderPrimaryBlocks Int
k [(Int, ByteString)]
inblocks
                    presentBlocks :: [Int]
presentBlocks = ((Int, ByteString) -> Int) -> [(Int, ByteString)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst [(Int, ByteString)]
inblocks'
                ForeignPtr CFEC -> (Ptr CFEC -> IO [ByteString]) -> IO [ByteString]
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr
                    ForeignPtr CFEC
params
                    ( \Ptr CFEC
cfec' -> do
                        [ByteString]
-> (Ptr (Ptr Word8) -> IO [ByteString]) -> IO [ByteString]
forall a. [ByteString] -> (Ptr (Ptr Word8) -> IO a) -> IO a
byteStringsToArray
                            (((Int, ByteString) -> ByteString)
-> [(Int, ByteString)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd [(Int, ByteString)]
inblocks')
                            ( \Ptr (Ptr Word8)
src -> do
                                [ByteString]
b <-
                                    Int -> Int -> (Ptr (Ptr Word8) -> IO ()) -> IO [ByteString]
createByteStringArray
                                        (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)
                                        Int
sz
                                        ( \Ptr (Ptr Word8)
out -> do
                                            [Int] -> (Ptr CUInt -> IO ()) -> IO ()
forall a. [Int] -> (Ptr CUInt -> IO a) -> IO a
uintCArray
                                                [Int]
presentBlocks
                                                ( \Ptr CUInt
block_nums -> do
                                                    Ptr CFEC
-> Ptr (Ptr Word8)
-> Ptr (Ptr Word8)
-> Ptr CUInt
-> CSize
-> IO ()
_decode Ptr CFEC
cfec' Ptr (Ptr Word8)
src Ptr (Ptr Word8)
out Ptr CUInt
block_nums (CSize -> IO ()) -> CSize -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz
                                                )
                                        )
                                let blocks :: [Int]
blocks = [Int
0 .. (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)] [Int] -> [Int] -> [Int]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Int]
presentBlocks
                                    tagged :: [(Int, ByteString)]
tagged = [Int] -> [ByteString] -> [(Int, ByteString)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
blocks [ByteString]
b
                                    allBlocks :: [(Int, ByteString)]
allBlocks = [(Int, ByteString)] -> [(Int, ByteString)]
forall a. [(Int, a)] -> [(Int, a)]
sortTagged ([(Int, ByteString)] -> [(Int, ByteString)])
-> [(Int, ByteString)] -> [(Int, ByteString)]
forall a b. (a -> b) -> a -> b
$ [(Int, ByteString)]
tagged [(Int, ByteString)] -> [(Int, ByteString)] -> [(Int, ByteString)]
forall a. [a] -> [a] -> [a]
++ [(Int, ByteString)]
inblocks'
                                [ByteString] -> IO [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> IO [ByteString])
-> [ByteString] -> IO [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take Int
k ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ ((Int, ByteString) -> ByteString)
-> [(Int, ByteString)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd [(Int, ByteString)]
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 ::
    -- | the number of parts requested
    Int ->
    -- | the data to be split
    B.ByteString ->
    IO [B.ByteString]
secureDivide :: Int -> ByteString -> IO [ByteString]
secureDivide Int
n ByteString
input
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> IO [ByteString]
forall a. HasCallStack => String -> a
error String
"secureDivide called with negative number of parts"
    | Bool
otherwise =
        String -> IOMode -> (Handle -> IO [ByteString]) -> IO [ByteString]
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile
            String
"/dev/urandom"
            IOMode
ReadMode
            ( \Handle
handle -> do
                let inner :: t -> ByteString -> IO [ByteString]
inner t
1 ByteString
bs = [ByteString] -> IO [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return [ByteString
bs]
                    inner t
n' ByteString
bs = do
                        ByteString
mask <- Handle -> Int -> IO ByteString
B.hGet Handle
handle (ByteString -> Int
B.length ByteString
bs)
                        let masked :: ByteString
masked = [Word8] -> ByteString
B.pack ([Word8] -> ByteString) -> [Word8] -> ByteString
forall a b. (a -> b) -> a -> b
$ (Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
B.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
bs ByteString
mask
                        [ByteString]
rest <- t -> ByteString -> IO [ByteString]
inner (t
n' t -> t -> t
forall a. Num a => a -> a -> a
- t
1) ByteString
masked
                        [ByteString] -> IO [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
mask ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
rest)
                Int -> ByteString -> IO [ByteString]
forall t. (Eq t, Num t) => t -> ByteString -> IO [ByteString]
inner Int
n ByteString
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 :: [ByteString] -> ByteString
secureCombine [] = String -> ByteString
forall a. HasCallStack => String -> a
error String
"Passed empty list of inputs to secureCombine"
secureCombine [ByteString
a] = ByteString
a
secureCombine [ByteString
a, ByteString
b] = [Word8] -> ByteString
B.pack ([Word8] -> ByteString) -> [Word8] -> ByteString
forall a b. (a -> b) -> a -> b
$ (Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
B.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
a ByteString
b
secureCombine (ByteString
a : [ByteString]
rest) = [Word8] -> ByteString
B.pack ([Word8] -> ByteString) -> [Word8] -> ByteString
forall a b. (a -> b) -> a -> b
$ (Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
B.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
a (ByteString -> [Word8]) -> ByteString -> [Word8]
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
secureCombine [ByteString]
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 ::
    -- | the number of blocks required to reconstruct
    Int ->
    -- | the total number of blocks
    Int ->
    -- | the data to divide
    B.ByteString ->
    -- | the resulting blocks
    [B.ByteString]
enFEC :: Int -> Int -> ByteString -> [ByteString]
enFEC Int
k Int
n ByteString
input = [ByteString]
taggedPrimaryBlocks [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString]
taggedSecondaryBlocks
  where
    taggedPrimaryBlocks :: [ByteString]
taggedPrimaryBlocks = (Word8 -> ByteString -> ByteString)
-> [Word8] -> [ByteString] -> [ByteString]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word8 -> ByteString -> ByteString
B.cons [Word8
0 ..] [ByteString]
primaryBlocks
    taggedSecondaryBlocks :: [ByteString]
taggedSecondaryBlocks = (Word8 -> ByteString -> ByteString)
-> [Word8] -> [ByteString] -> [ByteString]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word8 -> ByteString -> ByteString
B.cons [(Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) ..] [ByteString]
secondaryBlocks
    remainder :: Int
remainder = ByteString -> Int
B.length ByteString
input Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k
    paddingLength :: Int
paddingLength = if Int
remainder Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 then Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remainder else Int
k
    paddingBytes :: ByteString
paddingBytes = Int -> Word8 -> ByteString
B.replicate (Int
paddingLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word8
0 ByteString -> ByteString -> ByteString
`B.append` Word8 -> ByteString
B.singleton (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
paddingLength)
    divide :: Int -> ByteString -> [ByteString]
divide Int
a ByteString
bs
        | ByteString -> Bool
B.null ByteString
bs = []
        | Bool
otherwise = Int -> ByteString -> ByteString
B.take Int
a ByteString
bs ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Int -> ByteString -> [ByteString]
divide Int
a (Int -> ByteString -> ByteString
B.drop Int
a ByteString
bs)
    input' :: ByteString
input' = ByteString
input ByteString -> ByteString -> ByteString
`B.append` ByteString
paddingBytes
    blockSize :: Int
blockSize = ByteString -> Int
B.length ByteString
input' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
k
    primaryBlocks :: [ByteString]
primaryBlocks = Int -> ByteString -> [ByteString]
divide Int
blockSize ByteString
input'
    secondaryBlocks :: [ByteString]
secondaryBlocks = FECParams -> [ByteString] -> [ByteString]
encode FECParams
params [ByteString]
primaryBlocks
    params :: FECParams
params = Int -> Int -> FECParams
fec Int
k Int
n

-- | Reverses the operation of @enFEC@.
deFEC ::
    -- | the number of blocks required (matches call to @enFEC@)
    Int ->
    -- | the total number of blocks (matches call to @enFEC@)
    Int ->
    -- | a list of k, or more, blocks from @enFEC@
    [B.ByteString] ->
    B.ByteString
deFEC :: Int -> Int -> [ByteString] -> ByteString
deFEC Int
k Int
n [ByteString]
inputs
    | [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
inputs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k = String -> ByteString
forall a. HasCallStack => String -> a
error String
"Too few inputs to deFEC"
    | Bool
otherwise = Int -> ByteString -> ByteString
B.take (ByteString -> Int
B.length ByteString
fecOutput Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
paddingLength) ByteString
fecOutput
  where
    paddingLength :: Int
paddingLength = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ ByteString -> Word8
B.last ByteString
fecOutput
    inputs' :: [ByteString]
inputs' = Int -> [ByteString] -> [ByteString]
forall a. Int -> [a] -> [a]
take Int
k [ByteString]
inputs
    taggedInputs :: [(Int, ByteString)]
taggedInputs = (ByteString -> (Int, ByteString))
-> [ByteString] -> [(Int, ByteString)]
forall a b. (a -> b) -> [a] -> [b]
map (\ByteString
bs -> (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ ByteString -> Word8
B.head ByteString
bs, ByteString -> ByteString
B.tail ByteString
bs)) [ByteString]
inputs'
    fecOutput :: ByteString
fecOutput = [ByteString] -> ByteString
B.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ FECParams -> [(Int, ByteString)] -> [ByteString]
decode FECParams
params [(Int, ByteString)]
taggedInputs
    params :: FECParams
params = Int -> Int -> FECParams
fec Int
k Int
n