{-|
Module      : Botan.Low.ZFEC
Description : ZFEC Forward Error Correction
Copyright   : (c) 2023-2024, Apotheca Labs
              (c) 2024-2025, Haskell Foundation
License     : BSD-3-Clause
Maintainer  : joris@well-typed.com, 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.

-}

{-# LANGUAGE OverloadedStrings #-}

module Botan.Low.ZFEC (

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

  -- * ZFEC
    ZFECShare
  , zfecEncode
  , zfecDecode

  ) where

import           Botan.Bindings.ConstPtr (ConstPtr (..))
import           Botan.Bindings.ZFEC
import           Botan.Low.Error.Internal
import           Botan.Low.Internal.ByteString
import           Botan.Low.Make
import           Control.DeepSeq
import           Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import           Data.Word
import           Foreign.C.Types
import           Foreign.Marshal.Alloc
import           Foreign.Marshal.Array
import           Foreign.Ptr

{- $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 CInt -> IO ()
IO CInt -> IO ()
throwBotanIfNegative_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ CSize
-> CSize -> ConstPtr Word8 -> CSize -> Ptr (Ptr Word8) -> IO CInt
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 CInt -> IO ()
IO CInt -> IO ()
throwBotanIfNegative_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ CSize
-> CSize
-> ConstPtr CSize
-> ConstPtr (ConstPtr Word8)
-> CSize
-> Ptr (Ptr Word8)
-> IO CInt
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