{-|
Module      : Botan.Low.ZFEC
Description : ZFEC Forward Error Correction
Copyright   : (c) Leo D, 2023
License     : BSD-3-Clause
Maintainer  : leo@apotheca.io
Stability   : experimental
Portability : POSIX

Forward error correction takes an input and creates multiple
“shares”, such that any K of N shares is sufficient to recover
the entire original input.

-}

module Botan.Low.ZFEC
(

-- * Forward Error Correction
-- $introduction
-- * Usage
-- $usage

-- * ZFEC
  ZFECShare(..)
, zfecEncode
, zfecDecode

) where

import Control.Concurrent

import Data.Foldable

import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Unsafe as ByteString

import Botan.Bindings.ZFEC

import Botan.Low.Error
import Botan.Low.Make
import Botan.Low.Prelude

{- $introduction

The ZFEC module provides forward error correction compatible
with the zfec library.

Note

Specific to the ZFEC format, the first K generated shares are
identical to the original input data, followed by N-K shares of
error correcting code. This is very different from threshold
secret sharing, where having fewer than K shares gives no
information about the original input.

Warning

If a corrupted share is provided to the decoding algorithm, the
resulting decoding will be invalid. It is recommended to protect
shares using a technique such as a MAC or public key signature,
if corruption is likely in your application.

-}

{- $usage

Forward error correction takes an input and creates multiple
“shares”, such that any K of N shares is sufficient to recover
the entire original input.

First, we choose a K value appropriate to our message - the higher K is,
the smaller (but more numerous) the resulting shares will be:

> k = 7
> message = "The length of this message must be divisible by K"

> NOTE: ZFEC requires that the input length be exactly divisible by K; if
needed define a padding scheme to pad your input to the necessary
size.

We can calculate N = K + R, where R is the number of redundant shares,
meaning we can tolerate the loss of up to R shares and still recover
the original message.

We want 2 additional shares of redundancy, so we set R and N appropriately:

> r = 2
> n = k + r -- 7 + 2 = 9

Then, we encode the message into N shares:

> shares <- zfecEncode k n message
> length shares
> -- 9

Then, we can recover the message from any K of N shares:

> someShares <- take k <$> shuffle shares
> recoveredMessage <- zfecDecode k n someShares
> message == recoveredMessage
> -- True

-}

type ZFECShare = (Int, ByteString)

-- Or should this be?:
-- zfecEncode :: Int -> Int -> Int -> Input -> IO [ZFECShare]
-- zfecEncode k n shareSz input = ...
-- ^ is more 'raw'.

-- | Encode some bytes with certain ZFEC parameters.
--
-- NOTE: The length in bytes of input must be a multiple of K
zfecEncode
    :: Int              -- ^ __K__: the number of shares needed for recovery
    -> Int              -- ^ __N__: the number of shares generated
    -> ByteString       -- ^ __input__: the data to FEC
    -> IO [ZFECShare]   
zfecEncode :: Int -> Int -> ByteString -> IO [ZFECShare]
zfecEncode Int
k Int
n ByteString
input = ByteString
-> (Ptr Word8 -> CSize -> IO [ZFECShare]) -> IO [ZFECShare]
forall byte a. ByteString -> (Ptr byte -> CSize -> IO a) -> IO a
asBytesLen ByteString
input ((Ptr Word8 -> CSize -> IO [ZFECShare]) -> IO [ZFECShare])
-> (Ptr Word8 -> CSize -> IO [ZFECShare]) -> IO [ZFECShare]
forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
inputPtr CSize
inputLen -> do
    let shareSize :: Int
shareSize = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
inputLen) Int
k
    Int -> (Ptr Word8 -> IO [ZFECShare]) -> IO [ZFECShare]
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
shareSize) ((Ptr Word8 -> IO [ZFECShare]) -> IO [ZFECShare])
-> (Ptr Word8 -> IO [ZFECShare]) -> IO [ZFECShare]
forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
outPtr -> do
        Int -> (Ptr (Ptr Word8) -> IO [ZFECShare]) -> IO [ZFECShare]
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
n ((Ptr (Ptr Word8) -> IO [ZFECShare]) -> IO [ZFECShare])
-> (Ptr (Ptr Word8) -> IO [ZFECShare]) -> IO [ZFECShare]
forall a b. (a -> b) -> a -> b
$ \ (Ptr (Ptr Word8)
sharePtrArrayPtr :: Ptr (Ptr Word8)) -> do
            let sharePtrs :: [Ptr Word8]
sharePtrs = (Int -> Ptr Word8) -> [Int] -> [Ptr Word8]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Ptr Word8 -> Int -> Ptr Word8
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr Word8
outPtr (Int -> Ptr Word8) -> (Int -> Int) -> Int -> Ptr Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
shareSize)) [Int
0..(Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)]
            Ptr (Ptr Word8) -> [Ptr Word8] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr (Ptr Word8)
sharePtrArrayPtr [Ptr Word8]
sharePtrs
            HasCallStack => IO BotanErrorCode -> IO ()
IO BotanErrorCode -> IO ()
throwBotanIfNegative_ (IO BotanErrorCode -> IO ()) -> IO BotanErrorCode -> IO ()
forall a b. (a -> b) -> a -> b
$ CSize
-> CSize
-> ConstPtr Word8
-> CSize
-> Ptr (Ptr Word8)
-> IO BotanErrorCode
botan_zfec_encode
                (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k)
                (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
                (Ptr Word8 -> ConstPtr Word8
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr Word8
inputPtr)
                CSize
inputLen
                Ptr (Ptr Word8)
sharePtrArrayPtr
            [ByteString]
shares <- (Ptr Word8 -> IO ByteString) -> [Ptr Word8] -> IO [ByteString]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (CStringLen -> IO ByteString
ByteString.packCStringLen (CStringLen -> IO ByteString)
-> (Ptr Word8 -> CStringLen) -> Ptr Word8 -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,Int
shareSize) (Ptr CChar -> CStringLen)
-> (Ptr Word8 -> Ptr CChar) -> Ptr Word8 -> CStringLen
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Word8 -> Ptr CChar
forall a b. Ptr a -> Ptr b
castPtr) [Ptr Word8]
sharePtrs
            [ZFECShare] -> IO [ZFECShare]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([ZFECShare] -> IO [ZFECShare]) -> [ZFECShare] -> IO [ZFECShare]
forall a b. NFData a => (a -> b) -> a -> b
$!! [Int] -> [ByteString] -> [ZFECShare]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..(Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)] [ByteString]
shares

-- TODO: Throw a fit if shares are not equal length, not k shares
    
-- | Decode some previously encoded shares using certain ZFEC parameters.
--
-- NOTE: There must be at least K shares of equal length
zfecDecode
    :: Int              -- ^ __K__: the number of shares needed for recovery
    -> Int              -- ^ __N__: the total number of shares
    -> [ZFECShare]      -- ^ __inputs__: K previously encoded shares to decode
    -> IO ByteString    -- ^ __outputs__: An out parameter pointing to a fully allocated array of size
                        --   [N][size / K].  For all n in range, an encoded block will be
                        --   written to the memory starting at outputs[n][0].
zfecDecode :: Int -> Int -> [ZFECShare] -> IO ByteString
zfecDecode Int
_ Int
_ [] = ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
""
zfecDecode Int
k Int
n shares :: [ZFECShare]
shares@((Int
_,ByteString
share0):[ZFECShare]
_) = do
    Int -> (Ptr CSize -> IO ByteString) -> IO ByteString
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
k ((Ptr CSize -> IO ByteString) -> IO ByteString)
-> (Ptr CSize -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \ (Ptr CSize
indexesPtr :: Ptr CSize) -> do
        Ptr CSize -> [CSize] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr CSize
indexesPtr [CSize]
shareIndexes
        (forall a. ByteString -> (Ptr Word8 -> IO a) -> IO a)
-> [ByteString] -> ([Ptr Word8] -> IO ByteString) -> IO ByteString
forall typ ptr b.
(forall a. typ -> (ptr -> IO a) -> IO a)
-> [typ] -> ([ptr] -> IO b) -> IO b
withPtrs ByteString -> (Ptr Word8 -> IO a) -> IO a
forall a. ByteString -> (Ptr Word8 -> IO a) -> IO a
forall byte a. ByteString -> (Ptr byte -> IO a) -> IO a
unsafeAsBytes [ByteString]
shareBytes (([Ptr Word8] -> IO ByteString) -> IO ByteString)
-> ([Ptr Word8] -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \ ([Ptr Word8]
sharePtrs :: [Ptr Word8]) -> do
            Int -> (Ptr (Ptr Word8) -> IO ByteString) -> IO ByteString
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
k ((Ptr (Ptr Word8) -> IO ByteString) -> IO ByteString)
-> (Ptr (Ptr Word8) -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \ (Ptr (Ptr Word8)
sharePtrArrayPtr :: Ptr (Ptr Word8)) -> do
                Ptr (Ptr Word8) -> [Ptr Word8] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr (Ptr Word8)
sharePtrArrayPtr [Ptr Word8]
sharePtrs
                -- NOTE: This extra work may potentially be avoided by allocating a
                --  single contiguous block
                -- withPtrs (const $ allocaBytes shareSize) [0..(k-1)] $ \ outPtrs -> do
                --     allocaArray k $ \ outPtrArrayPtr -> do
                --         pokeArray outPtrArrayPtr outPtrs
                --         throwBotanIfNegative_ $ botan_zfec_decode
                --             (fromIntegral k)
                --             (fromIntegral n)
                --             indexesPtr
                --             sharePtrArrayPtr
                --             (fromIntegral shareSize)
                --             outPtrArrayPtr
                --         decodedShares <- traverse (ByteString.unsafePackCStringLen . (,shareSize) . castPtr) outPtrs
                --         return $!! ByteString.copy $ ByteString.concat decodedShares
                -- Single contiguous block method
                -- This way is probably superior absent any surprise alignment issues
                Int -> (Ptr Word8 -> IO ()) -> IO ByteString
forall byte. Int -> (Ptr byte -> IO ()) -> IO ByteString
allocBytes (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
shareSize) ((Ptr Word8 -> IO ()) -> IO ByteString)
-> (Ptr Word8 -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
outPtr -> do
                    Int -> (Ptr (Ptr Word8) -> IO ()) -> IO ()
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
n ((Ptr (Ptr Word8) -> IO ()) -> IO ())
-> (Ptr (Ptr Word8) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ (Ptr (Ptr Word8)
outPtrArrayPtr :: Ptr (Ptr Word8)) -> do
                        let outPtrs :: [Ptr Word8]
outPtrs = (Int -> Ptr Word8) -> [Int] -> [Ptr Word8]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Ptr Word8 -> Int -> Ptr Word8
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr Word8
outPtr (Int -> Ptr Word8) -> (Int -> Int) -> Int -> Ptr Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
shareSize)) [Int
0..(Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)]
                        Ptr (Ptr Word8) -> [Ptr Word8] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr (Ptr Word8)
outPtrArrayPtr [Ptr Word8]
outPtrs
                        HasCallStack => IO BotanErrorCode -> IO ()
IO BotanErrorCode -> IO ()
throwBotanIfNegative_ (IO BotanErrorCode -> IO ()) -> IO BotanErrorCode -> IO ()
forall a b. (a -> b) -> a -> b
$ CSize
-> CSize
-> ConstPtr CSize
-> ConstPtr (ConstPtr Word8)
-> CSize
-> Ptr (Ptr Word8)
-> IO BotanErrorCode
botan_zfec_decode
                            (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k)
                            (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
                            (Ptr CSize -> ConstPtr CSize
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr CSize
indexesPtr)
                            -- NOTE: Use of castPtr here because allocating
                            --  as a ConstPtr (ConstPtr a) is tedious
                            (Ptr (ConstPtr Word8) -> ConstPtr (ConstPtr Word8)
forall a. Ptr a -> ConstPtr a
ConstPtr (Ptr (ConstPtr Word8) -> ConstPtr (ConstPtr Word8))
-> Ptr (ConstPtr Word8) -> ConstPtr (ConstPtr Word8)
forall a b. (a -> b) -> a -> b
$ Ptr (Ptr Word8) -> Ptr (ConstPtr Word8)
forall a b. Ptr a -> Ptr b
castPtr Ptr (Ptr Word8)
sharePtrArrayPtr)
                            (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
shareSize)
                            Ptr (Ptr Word8)
outPtrArrayPtr
    where
        shareIndexes :: [CSize]
shareIndexes = (ZFECShare -> CSize) -> [ZFECShare] -> [CSize]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CSize) -> (ZFECShare -> Int) -> ZFECShare -> CSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ZFECShare -> Int
forall a b. (a, b) -> a
fst) [ZFECShare]
shares
        shareBytes :: [ByteString]
shareBytes = (ZFECShare -> ByteString) -> [ZFECShare] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ZFECShare -> ByteString
forall a b. (a, b) -> b
snd [ZFECShare]
shares
        shareSize :: Int
shareSize = ByteString -> Int
ByteString.length ByteString
share0