FEC.hs

{-# 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
    { _cfec :: !(ForeignPtr CFEC)
    , paramK :: Int
    , paramN :: Int
    }
    deriving (Generic)

-- Provide an NFData instance so it's possible to use a FECParams in a
-- Criterion benchmark.
instance NFData FECParams where
    rnf FECParams{_cfec, paramK, 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?
        rnf paramK `seq` rnf paramN

instance Show FECParams where
    show (FECParams _ k n) = "FEC (" ++ show k ++ ", " ++ show n ++ ")"

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 k n
    | k > n = False
    | k < 1 = False
    | n < 1 = False
    | n > 255 = False
    | otherwise = True

{- | The underlying library signaled that it has not been properly initialized
 yet.  Use @initialize@ to initialize it.
-}
data Uninitialized = Uninitialized deriving (Ord, Eq, 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 = unsafePerformIO newLock

-- | Initialize the library.  This must be done before other APIs can succeed.
initialize :: IO ()
initialize = withLock _initializationLock _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 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)
                    -- new will return null if the library hasn't been
                    -- initialized.
                    if cfec' == nullPtr
                        then throwIO Uninitialized
                        else do
                            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 = withArray . map 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 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) . 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 ::
    -- | the number of buffers requested
    Int ->
    -- | the size of each buffer
    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)]
                )
        )

{- | 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 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 ->
    -- | a list of @k@ blocks and their index
    [(Int, B.ByteString)] ->
    -- | a list the @k@ primary blocks
    [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) . 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 ::
    -- | the number of parts requested
    Int ->
    -- | the data to be split
    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
            )

{- | 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 ::
    -- | 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 k n input = taggedPrimaryBlocks ++ taggedSecondaryBlocks
  where
    taggedPrimaryBlocks = zipWith B.cons [0 ..] primaryBlocks
    taggedSecondaryBlocks = zipWith B.cons [(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 ::
    -- | 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 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