ZFEC.hs

{-|
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 k n input = asBytesLen input $ \ inputPtr inputLen -> do
    let shareSize = div (fromIntegral inputLen) k
    allocaBytes (n * shareSize) $ \ outPtr -> do
        allocaArray n $ \ (sharePtrArrayPtr :: Ptr (Ptr Word8)) -> do
            let sharePtrs = fmap (advancePtr outPtr . (* shareSize)) [0..(n-1)]
            pokeArray sharePtrArrayPtr sharePtrs
            throwBotanIfNegative_ $ botan_zfec_encode
                (fromIntegral k)
                (fromIntegral n)
                (ConstPtr inputPtr)
                inputLen
                sharePtrArrayPtr
            shares <- traverse (ByteString.packCStringLen . (,shareSize) . castPtr) sharePtrs
            return $!! zip [0..(n-1)] 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 _ _ [] = return ""
zfecDecode k n shares@((_,share0):_) = do
    allocaArray k $ \ (indexesPtr :: Ptr CSize) -> do
        pokeArray indexesPtr shareIndexes
        withPtrs unsafeAsBytes shareBytes $ \ (sharePtrs :: [Ptr Word8]) -> do
            allocaArray k $ \ (sharePtrArrayPtr :: Ptr (Ptr Word8)) -> do
                pokeArray sharePtrArrayPtr 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
                allocBytes (k * shareSize) $ \ outPtr -> do
                    allocaArray n $ \ (outPtrArrayPtr :: Ptr (Ptr Word8)) -> do
                        let outPtrs = fmap (advancePtr outPtr . (* shareSize)) [0..(n-1)]
                        pokeArray outPtrArrayPtr outPtrs
                        throwBotanIfNegative_ $ botan_zfec_decode
                            (fromIntegral k)
                            (fromIntegral n)
                            (ConstPtr indexesPtr)
                            -- NOTE: Use of castPtr here because allocating
                            --  as a ConstPtr (ConstPtr a) is tedious
                            (ConstPtr $ castPtr sharePtrArrayPtr)
                            (fromIntegral shareSize)
                            outPtrArrayPtr
    where
        shareIndexes = fmap (fromIntegral . fst) shares
        shareBytes = fmap snd shares
        shareSize = ByteString.length share0