{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}

--
-- Glossary
--
-- CHK - An encryption and encoding scheme for storing immutable data.
--
-- Data - The plaintext used to construct CHK.
--
-- Share - One complete unit of encrypted and FEC encoded data.
--
-- Segment - One piece of ciphertext used to construct CHK.  All segments
--     belonging to a CHK are the same size except the last one may be short.
--
-- k, required - The number of "primary" erasure-encoding outputs.  Equal to
--     the minimum number of erasure-encoding outputs needed to reconstruct
--     the erasure-encoding input.
--
-- n, total - The total number of erasure-encoding outputs.  Always greater
--     than or equal to required.
--
-- Block - One output resulting from erasure-encoding one segment using
--     required, total.  If necessary, the input segment is nul-padded so its
--     size is a multiple of required.
--
-- Plaintext Hash Tree - Not actually implemented by Tahoe so I'm not really
--     sure.  Probably something like a sha256d merkle tree where the leaves
--     are hashes of the plaintext corresponding to each ciphertext segment.
--     Since all shares for a CHK are derived from the same plaintext, every
--     share has the same plaintext hash tree.
--
-- Crypttext Hash Tree - A sha256d merkle tree where the leaves are hashes of
--     the ciphertext segments.  Since all shares for a CHK are derived from
--     the same ciphertext, every share has the same ciphertext hash tree.
--
-- Crypttext Root Hash - The hash at the root of Crypttext Hash Tree.
--
-- Block Hash Tree - A sha256d merkle tree where the leaves are hashes of the
--     blocks.  Since the erasure-encoding output is different for each share,
--     every share has a different block hash tree.
--
-- Share Hash Tree - A sha256d merkle tree where the leaves are the root
--     hashes of the block hash trees for all shares.
--
-- Share Hashes - A list of hashes from the Share Hash Tree which are required
--     to verify one block hash tree.  Each share includes the Share Hashes
--     required to verify the Block Hash Tree contained within that share.
--     Since every share contains a different Block Hash Tree, every share
--     contains a different list of Share Hashes.  Each Share Hash in this
--     list is accompanied by information about its position in the Share Hash
--     Tree though it may not be strictly required (since it could be inferred
--     from position in the list).
--
-- URI Extension - A collection of metadata describing the encryption and
--     encoding used to create the CHK and (largely) necessary for either
--     decoding or verifying the integrity of the contained data.

module Tahoe.CHK (
    zfec,
    zunfec,
    encode,
    decode,
    padCiphertext,
    segmentCiphertext,
) where

import qualified Codec.FEC as ZFEC
import Crypto.Cipher.AES128 (
    AESKey128,
 )
import Data.Int (Int64)
import Data.Word (Word64)

-- import Debug.Trace

import Crypto.Hash (
    Context,
    HashAlgorithm (hashDigestSize),
    SHA256 (SHA256),
    hashFinalize,
    hashInit,
    hashUpdate,
 )
import Data.Bifunctor (first, second)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Foldable (foldlM)
import Data.List (sort, transpose)
import Data.List.Extra (
    snoc,
 )
import Data.Maybe (fromJust)
import qualified Tahoe.CHK.Capability as Cap
import Tahoe.CHK.Crypto (
    blockHash,
    ciphertextSegmentHash,
    ciphertextTag,
    sha256,
    uriExtensionHash,
 )
import Tahoe.CHK.Merkle (
    MerkleTree,
    buildTreeOutOfAllTheNodes,
    leafHashes,
    leafNumberToNodeNumber,
    makeTreePartial,
    neededHashes,
    rootHash,
 )
import Tahoe.CHK.Share (Share (..))
import Tahoe.CHK.Types (
    BlockHash,
    CrypttextHash,
    Parameters (..),
    Required,
    requiredToInt,
    totalToInt,
 )
import Tahoe.CHK.URIExtension (
    URIExtension (..),
 )
import Tahoe.Netstring (
    netstring,
 )
import Tahoe.Util (
    ceilDiv,
    chunkedBy,
    nextMultipleOf,
    nextPowerOf,
 )

-- | Erasure encode some bytes using ZFEC.
zfec ::
    -- | The number of outputs that will be required to reverse the encoding.
    -- Also known as `k`.
    Int ->
    -- | The total number of outputs to produce.  Also known as `n`.
    Int ->
    -- | Application data to divide into encoding inputs.
    B.ByteString ->
    -- | `n` encoding outputs.
    IO [B.ByteString]
zfec :: Int -> Int -> ByteString -> IO [ByteString]
zfec Int
k Int
n ByteString
segment =
    [ByteString] -> IO [ByteString]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ByteString] -> IO [ByteString])
-> [ByteString] -> IO [ByteString]
forall a b. (a -> b) -> a -> b
$ [ByteString]
chunks [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ FECParams -> [ByteString] -> [ByteString]
ZFEC.encode (Int -> Int -> FECParams
ZFEC.fec Int
k Int
n) [ByteString]
chunks
  where
    chunks_ :: [ByteString]
chunks_ = Int -> ByteString -> [ByteString]
chunkedBy (ByteString -> Int
B.length ByteString
segment Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
k) ByteString
segment
    _msg :: String
_msg =
        String
"zfec"
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" k="
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
k
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" n="
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", segment len "
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
B.length ByteString
segment)
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", chunk lengths "
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [Int] -> String
forall a. Show a => a -> String
show ((ByteString -> Int) -> [ByteString] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ByteString -> Int
B.length [ByteString]
chunks_)
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", segment "
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
segment
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"-> chunks "
            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [ByteString] -> String
forall a. Show a => a -> String
show [ByteString]
chunks_
    chunks :: [ByteString]
chunks = {- trace _msg -} [ByteString]
chunks_

-- | Version of `zfec` that operates on lazy ByteStrings.
zfecLazy :: Int -> Int -> LB.ByteString -> IO [LB.ByteString]
zfecLazy :: Int -> Int -> ByteString -> IO [ByteString]
zfecLazy Int
k Int
n ByteString
segment = (ByteString -> ByteString
LB.fromStrict (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) ([ByteString] -> [ByteString])
-> IO [ByteString] -> IO [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> ByteString -> IO [ByteString]
zfec Int
k Int
n (ByteString -> ByteString
LB.toStrict ByteString
segment)

-- | Erasure decode some bytes using ZFEC.
zunfec ::
    -- | The `k` parameter used when encoding the data to decode.
    Int ->
    -- | The `n` parameter used when encoding the data to decode.
    Int ->
    -- | The encoding outputs annotated with their position (or "share number").
    [(Int, B.ByteString)] ->
    -- | The bytes which were originally encoded.
    IO B.ByteString
zunfec :: Int -> Int -> [(Int, ByteString)] -> IO ByteString
zunfec Int
k Int
n [(Int, ByteString)]
blocks = ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
B.concat (FECParams -> [(Int, ByteString)] -> [ByteString]
ZFEC.decode (Int -> Int -> FECParams
ZFEC.fec Int
k Int
n) [(Int, ByteString)]
blocks)

-- | Version of `zunfec` that operates on lazy ByteStrings.
zunfecLazy :: Int -> Int -> [(Int, LB.ByteString)] -> IO LB.ByteString
zunfecLazy :: Int -> Int -> [(Int, ByteString)] -> IO ByteString
zunfecLazy Int
k Int
n [(Int, ByteString)]
blocks = do
    ByteString
segment_ <- ByteString -> ByteString
LB.fromStrict (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [(Int, ByteString)] -> IO ByteString
zunfec Int
k Int
n ((ByteString -> ByteString)
-> (Int, ByteString) -> (Int, ByteString)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ByteString -> ByteString
LB.toStrict ((Int, ByteString) -> (Int, ByteString))
-> [(Int, ByteString)] -> [(Int, ByteString)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, ByteString)]
blocks)
    let _msg :: String
_msg =
            String
"zunfec"
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" k="
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
k
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" n="
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" blocks="
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [(Int, ByteString)] -> String
forall a. Show a => a -> String
show [(Int, ByteString)]
blocks
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" -> segment "
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
segment_

    -- pure (trace _msg) segment_
    ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
segment_

{- | Represent progress encoding some ciphertext into a CHK share.  This
 carries along intermediate hash values used at the end to build extra
 self-authenticating fields into the share.
-}
data EncodingState = CPState
    { -- A single hash of all crypttext segments encoded so far.
      EncodingState -> Context SHA256
cpCrypttextHash :: Crypto.Hash.Context Crypto.Hash.SHA256
    , -- A list of hashes of each ciphertext segment encoded so far
      EncodingState -> [ByteString]
cpCrypttextHashes :: [CrypttextHash]
    , -- Hashes of blocks encoded so far.
      EncodingState -> [[ByteString]]
cpBlockHashes :: [[BlockHash]]
    , -- Blocks encoded so far.
      EncodingState -> [[ByteString]]
cpBlocks :: [[LB.ByteString]]
    }

-- | The initial state for CHK encoding.
initEncodingState :: EncodingState
initEncodingState :: EncodingState
initEncodingState =
    CPState :: Context SHA256
-> [ByteString]
-> [[ByteString]]
-> [[ByteString]]
-> EncodingState
CPState
        { cpCrypttextHash :: Context SHA256
cpCrypttextHash = Context SHA256 -> ByteString -> Context SHA256
forall ba a.
(ByteArrayAccess ba, HashAlgorithm a) =>
Context a -> ba -> Context a
hashUpdate (Context SHA256
forall a. HashAlgorithm a => Context a
hashInit :: Context SHA256) (ByteString -> ByteString
netstring ByteString
ciphertextTag)
        , cpCrypttextHashes :: [ByteString]
cpCrypttextHashes = [ByteString]
forall a. Monoid a => a
mempty
        , cpBlockHashes :: [[ByteString]]
cpBlockHashes = [[ByteString]]
forall a. Monoid a => a
mempty
        , cpBlocks :: [[ByteString]]
cpBlocks = [[ByteString]]
forall a. Monoid a => a
mempty
        }

{- | Split a full ciphertext string into the separate ciphertext segments
 required by most of CHK encoding.
-}
segmentCiphertext ::
    -- | The encoding parameters which determine how to split the ciphertext.
    Parameters ->
    -- | The ciphertext.
    LB.ByteString ->
    -- | The segments.
    [LB.ByteString]
segmentCiphertext :: Parameters -> ByteString -> [ByteString]
segmentCiphertext Parameters{SegmentSize
paramSegmentSize :: Parameters -> SegmentSize
paramSegmentSize :: SegmentSize
paramSegmentSize} ByteString
ciphertext =
    [ByteString]
result
  where
    result :: [ByteString]
result = {- trace ("segmentCiphertext: " <> show ciphertext) -} [ByteString]
result_
    result_ :: [ByteString]
result_ = ByteString -> ByteString
LB.fromStrict (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ByteString -> [ByteString]
chunkedBy (SegmentSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral SegmentSize
paramSegmentSize) (ByteString -> ByteString
LB.toStrict ByteString
ciphertext)

{- | Process ciphertext into blocks, carrying hashes computed along the way as
 state.
-}
processCiphertext :: Parameters -> [LB.ByteString] -> IO EncodingState
processCiphertext :: Parameters -> [ByteString] -> IO EncodingState
processCiphertext Parameters{Required
paramRequiredShares :: Parameters -> Required
paramRequiredShares :: Required
paramRequiredShares, Required
paramTotalShares :: Parameters -> Required
paramTotalShares :: Required
paramTotalShares} =
    (EncodingState -> ByteString -> IO EncodingState)
-> EncodingState -> [ByteString] -> IO EncodingState
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM EncodingState -> ByteString -> IO EncodingState
processSegment EncodingState
initEncodingState
  where
    processSegment :: EncodingState -> ByteString -> IO EncodingState
processSegment CPState{[[ByteString]]
[[ByteString]]
[ByteString]
Context SHA256
cpBlocks :: [[ByteString]]
cpBlockHashes :: [[ByteString]]
cpCrypttextHashes :: [ByteString]
cpCrypttextHash :: Context SHA256
cpBlocks :: EncodingState -> [[ByteString]]
cpBlockHashes :: EncodingState -> [[ByteString]]
cpCrypttextHashes :: EncodingState -> [ByteString]
cpCrypttextHash :: EncodingState -> Context SHA256
..} ByteString
segment = do
        -- Produce the FEC blocks for this piece of ciphertext.
        [ByteString]
blocks <-
            Int -> Int -> ByteString -> IO [ByteString]
zfecLazy
                (Required -> Int
requiredToInt Required
paramRequiredShares)
                (Required -> Int
totalToInt Required
paramTotalShares)
                (Required -> ByteString -> ByteString
padCiphertext Required
paramRequiredShares ByteString
segment)
        EncodingState -> IO EncodingState
forall (f :: * -> *) a. Applicative f => a -> f a
pure (EncodingState -> IO EncodingState)
-> EncodingState -> IO EncodingState
forall a b. (a -> b) -> a -> b
$
            CPState :: Context SHA256
-> [ByteString]
-> [[ByteString]]
-> [[ByteString]]
-> EncodingState
CPState
                { cpCrypttextHash :: Context SHA256
cpCrypttextHash = Context SHA256 -> ByteString -> Context SHA256
forall ba a.
(ByteArrayAccess ba, HashAlgorithm a) =>
Context a -> ba -> Context a
hashUpdate Context SHA256
cpCrypttextHash (ByteString -> ByteString
LB.toStrict ByteString
segment)
                , cpCrypttextHashes :: [ByteString]
cpCrypttextHashes = [ByteString] -> ByteString -> [ByteString]
forall a. [a] -> a -> [a]
snoc [ByteString]
cpCrypttextHashes (ByteString -> ByteString
ciphertextSegmentHash (ByteString -> ByteString
LB.toStrict ByteString
segment))
                , cpBlockHashes :: [[ByteString]]
cpBlockHashes = [[ByteString]] -> [ByteString] -> [[ByteString]]
forall a. [a] -> a -> [a]
snoc [[ByteString]]
cpBlockHashes (ByteString -> ByteString
blockHash (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LB.toStrict (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ByteString]
blocks)
                , cpBlocks :: [[ByteString]]
cpBlocks = [[ByteString]] -> [ByteString] -> [[ByteString]]
forall a. [a] -> a -> [a]
snoc [[ByteString]]
cpBlocks [ByteString]
blocks
                }

-- Compute the correctly padded ciphertext.  The only ciphertext which is
-- expected to require padding is the final segment - in case the original
-- ciphertext did not have a length that was a multiple of the `required`
-- parameter.
--
-- allmydata.immutable.encode.Encoder._gather_data NUL pads up to num_chunks
-- times input_chunk_size.  num_chunks is our requiredShares.
-- input_chunk_size is taken from codec.get_block_size() which returns
-- codec.share_size.  share_size is div_ceil(data_size, required_shares).
-- data_size is our segmentSize and required_shares is our requiredShares.
padCiphertext :: Required -> LB.ByteString -> LB.ByteString
padCiphertext :: Required -> ByteString -> ByteString
padCiphertext Required
requiredShares ByteString
bs
    | Int64
paddingLength Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
0 = ByteString
bs ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int64 -> Word8 -> ByteString
LB.replicate Int64
paddingLength Word8
0x00
    | Bool
otherwise = ByteString
bs
  where
    desiredLength :: Int64
desiredLength = Required -> Int64 -> Int64
forall m v. (Integral m, Integral v) => m -> v -> v
nextMultipleOf Required
requiredShares (ByteString -> Int64
LB.length ByteString
bs)
    paddingLength :: Int64
paddingLength = Int64
desiredLength Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- ByteString -> Int64
LB.length ByteString
bs

{- | Encode some application data (typically ciphertext, but this function only
 weakly assumes this is the case) into some CHK shares.

 This replaces much of allmydata.immutable.encode.
-}
encode ::
    -- | The encryption/decryption key.
    AESKey128 ->
    -- | The ZFEC parameters for this encoding.  This determines how many shares
    -- will come out of this function.
    Parameters ->
    -- | The data to encode.  This is typically ciphertext.
    LB.ByteString ->
    -- | An IO which can be evaluated to get the encoded share data and the
    -- read capability.  The number of Shares will equal the `total` value
    -- from the given Parameters.
    IO ([Share], Cap.Reader)
encode :: AESKey128 -> Parameters -> ByteString -> IO ([Share], Reader)
encode AESKey128
readKey initParams :: Parameters
initParams@(Parameters SegmentSize
maximumSegmentSize Required
total Word8
_ Required
required) ByteString
ciphertext =
    Parameters -> [ByteString] -> IO EncodingState
processCiphertext Parameters
p (Parameters -> ByteString -> [ByteString]
segmentCiphertext Parameters
p ByteString
ciphertext) IO EncodingState
-> (EncodingState -> IO ([Share], Reader)) -> IO ([Share], Reader)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \CPState{[[ByteString]]
[[ByteString]]
[ByteString]
Context SHA256
cpBlocks :: [[ByteString]]
cpBlockHashes :: [[ByteString]]
cpCrypttextHashes :: [ByteString]
cpCrypttextHash :: Context SHA256
cpBlocks :: EncodingState -> [[ByteString]]
cpBlockHashes :: EncodingState -> [[ByteString]]
cpCrypttextHashes :: EncodingState -> [ByteString]
cpCrypttextHash :: EncodingState -> Context SHA256
..} ->
        let -- The number of segments encoded in the share.  There are the same number
            -- of plaintext and ciphertext segments and this is also the number of
            -- blocks in each share (though each share may have a different _value_
            -- for each block).
            --
            -- allmydata.immutable.encode.Encoder._got_all_encoding_parameters
            numSegments :: Int
numSegments = [[ByteString]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[ByteString]]
cpBlocks

            -- Our merkle trees need a number of leaves equal to a power of 2.
            -- Compute that here so we can pad as necessary.
            --
            -- allmydata.immutable.layout.WriteBucketProxy
            effectiveSegments :: Int
effectiveSegments = Int -> Int -> Int
forall p. (Ord p, Num p) => p -> p -> p
nextPowerOf Int
2 Int
numSegments

            -- XXX Unused by Tahoe so we don't even try for a sensible value right
            -- now.  Just fill it with zeros.
            --
            -- As long as we calculate a valid number of nodes for a tree
            -- buildTreeOutOfAllTheNodes won't give us a Nothing back ... cross
            -- your fingers.
            Just MerkleTree
plaintextHashTree =
                [ByteString] -> Maybe MerkleTree
buildTreeOutOfAllTheNodes
                    -- We have to fill the *whole* tree with nul, not just the
                    -- leaves.  Compute the total number of nodes in a tree that
                    -- can hold our number of segments.
                    ([ByteString] -> Maybe MerkleTree)
-> (ByteString -> [ByteString]) -> ByteString -> Maybe MerkleTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> [ByteString]
forall a. Int -> a -> [a]
replicate (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
effectiveSegments Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    -- And make every node all nul.
                    (ByteString -> Maybe MerkleTree) -> ByteString -> Maybe MerkleTree
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString
B.replicate (SHA256 -> Int
forall a. HashAlgorithm a => a -> Int
hashDigestSize SHA256
SHA256) Word8
0

            -- The merkle tree of ciphertext segment hashes.
            crypttextHashTree :: MerkleTree
crypttextHashTree = [ByteString] -> MerkleTree
makeTreePartial [ByteString]
cpCrypttextHashes

            -- shareTree is a MerkleTree of MerkleTree
            shareTree :: MerkleTree
shareTree =
                -- trace ("shareTree: " <> show shareTree')
                MerkleTree
shareTree'
              where
                shareTree' :: MerkleTree
shareTree' = [MerkleTree] -> MerkleTree
makeShareTree ([MerkleTree] -> MerkleTree)
-> ([[ByteString]] -> [MerkleTree]) -> [[ByteString]] -> MerkleTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([ByteString] -> MerkleTree) -> [[ByteString]] -> [MerkleTree]
forall a b. (a -> b) -> [a] -> [b]
map [ByteString] -> MerkleTree
makeTreePartial ([[ByteString]] -> [MerkleTree])
-> ([[ByteString]] -> [[ByteString]])
-> [[ByteString]]
-> [MerkleTree]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[ByteString]] -> [[ByteString]]
forall a. [[a]] -> [[a]]
transpose ([[ByteString]] -> MerkleTree) -> [[ByteString]] -> MerkleTree
forall a b. (a -> b) -> a -> b
$ [[ByteString]]
cpBlockHashes

            -- A bag of additional metadata about the share and encoded object.
            uriExtension :: URIExtension
uriExtension =
                URIExtension :: ByteString
-> Parameters
-> Parameters
-> SegmentSize
-> SegmentSize
-> Int
-> Required
-> Required
-> ByteString
-> ByteString
-> ByteString
-> URIExtension
URIExtension
                    { uriExtCodecName :: ByteString
uriExtCodecName = ByteString
"crs"
                    , uriExtCodecParams :: Parameters
uriExtCodecParams = Parameters
p -- trace ("Params: " <> show p) p
                    , uriExtSize :: SegmentSize
uriExtSize = Int64 -> SegmentSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> SegmentSize) -> Int64 -> SegmentSize
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
LB.length ByteString
ciphertext
                    , uriExtSegmentSize :: SegmentSize
uriExtSegmentSize = SegmentSize
segmentSize
                    , uriExtNeededShares :: Required
uriExtNeededShares = Required
required
                    , uriExtTotalShares :: Required
uriExtTotalShares = Required
total
                    , uriExtNumSegments :: Int
uriExtNumSegments = Int
numSegments
                    , uriExtTailCodecParams :: Parameters
uriExtTailCodecParams = Parameters -> Int64 -> Parameters
forall a. Integral a => Parameters -> a -> Parameters
tailParams Parameters
p (ByteString -> Int64
LB.length ByteString
ciphertext)
                    , uriExtCrypttextHash :: ByteString
uriExtCrypttextHash = Context SHA256 -> ByteString
makeCrypttextHash Context SHA256
cpCrypttextHash
                    , uriExtCrypttextRootHash :: ByteString
uriExtCrypttextRootHash = [ByteString] -> ByteString
makeCrypttextRootHash [ByteString]
cpCrypttextHashes
                    , uriExtShareRootHash :: ByteString
uriExtShareRootHash = MerkleTree -> ByteString
rootHash MerkleTree
shareTree
                    }

            -- The read capability for the encoded object.
            cap :: Reader
cap =
                AESKey128
-> ByteString -> Required -> Required -> SegmentSize -> Reader
Cap.makeReader
                    AESKey128
readKey
                    (URIExtension -> ByteString
uriExtensionHash URIExtension
uriExtension)
                    Required
required
                    Required
total
                    (Int64 -> SegmentSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> SegmentSize) -> Int64 -> SegmentSize
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
LB.length ByteString
ciphertext)

            toShare :: Int -> [ByteString] -> [ByteString] -> Share
toShare Int
sharenum [ByteString]
blocks [ByteString]
blockHashes =
                Share :: Word64
-> Word64
-> [ByteString]
-> MerkleTree
-> MerkleTree
-> MerkleTree
-> [(Word8, ByteString)]
-> URIExtension
-> Share
Share
                    { shareBlockSize :: Word64
shareBlockSize = Parameters -> Word64
shareBlockSize Parameters
p
                    , shareDataSize :: Word64
shareDataSize = Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Word64) -> Int64 -> Word64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
LB.length ByteString
ciphertext Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`ceilDiv` Required -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
required
                    , shareBlocks :: [ByteString]
shareBlocks = [ByteString]
blocks
                    , sharePlaintextHashTree :: MerkleTree
sharePlaintextHashTree = MerkleTree
plaintextHashTree
                    , shareCrypttextHashTree :: MerkleTree
shareCrypttextHashTree = MerkleTree
crypttextHashTree
                    , shareBlockHashTree :: MerkleTree
shareBlockHashTree = [ByteString] -> MerkleTree
makeTreePartial [ByteString]
blockHashes
                    , shareNeededHashes :: [(Word8, ByteString)]
shareNeededHashes = [(Word8, ByteString)] -> [(Word8, ByteString)]
forall a. Ord a => [a] -> [a]
sort ([(Word8, ByteString)] -> [(Word8, ByteString)])
-> ([(Int, ByteString)] -> [(Word8, ByteString)])
-> [(Int, ByteString)]
-> [(Word8, ByteString)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, ByteString) -> (Word8, ByteString))
-> [(Int, ByteString)] -> [(Word8, ByteString)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> Word8) -> (Int, ByteString) -> (Word8, ByteString)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ([(Int, ByteString)] -> [(Word8, ByteString)])
-> [(Int, ByteString)] -> [(Word8, ByteString)]
forall a b. (a -> b) -> a -> b
$ MerkleTree -> Int -> [(Int, ByteString)]
computeNeededShares MerkleTree
shareTree Int
sharenum
                    , shareURIExtension :: URIExtension
shareURIExtension = URIExtension
uriExtension
                    }

            -- The size in bytes of one erasure-encoded block of data.
            -- allmydata.immutable.encode.Encoder._got_all_encoding_parameters +
            -- allmydata.codec.CRSEncoder.set_params
            shareBlockSize :: Parameters -> Word64
            shareBlockSize :: Parameters -> Word64
shareBlockSize Parameters{SegmentSize
paramSegmentSize :: SegmentSize
paramSegmentSize :: Parameters -> SegmentSize
paramSegmentSize, Required
paramRequiredShares :: Required
paramRequiredShares :: Parameters -> Required
paramRequiredShares} =
                SegmentSize -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral SegmentSize
paramSegmentSize Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`ceilDiv` Required -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
paramRequiredShares
         in ([Share], Reader) -> IO ([Share], Reader)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( (Int -> [ByteString] -> [ByteString] -> Share)
-> [Int] -> [[ByteString]] -> [[ByteString]] -> [Share]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> [ByteString] -> [ByteString] -> Share
toShare [Int
0 ..] ([[ByteString]] -> [[ByteString]]
forall a. [[a]] -> [[a]]
transpose [[ByteString]]
cpBlocks) ([[ByteString]] -> [[ByteString]]
forall a. [[a]] -> [[a]]
transpose [[ByteString]]
cpBlockHashes)
                , Reader
cap
                )
  where
    -- If we have little enough ciphertext, the maximum configured segment
    -- size may be greater than the length of the single segment we produce.
    -- Segment size is also required to be a multiple of the number of
    -- required shares so that segments can be evenly divided across the
    -- shares.
    p :: Parameters
p@(Parameters SegmentSize
segmentSize Required
_ Word8
_ Required
required') =
        Parameters
initParams
            { paramSegmentSize :: SegmentSize
paramSegmentSize = Required -> SegmentSize -> SegmentSize
forall m v. (Integral m, Integral v) => m -> v -> v
nextMultipleOf Required
required' (SegmentSize -> SegmentSize) -> SegmentSize -> SegmentSize
forall a b. (a -> b) -> a -> b
$ SegmentSize -> SegmentSize -> SegmentSize
forall a. Ord a => a -> a -> a
min SegmentSize
maximumSegmentSize (Int64 -> SegmentSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> SegmentSize) -> Int64 -> SegmentSize
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
LB.length ByteString
ciphertext)
            }

{- | Decode some CHK shares to recover some application data.  This is roughly
 the inverse of ``encode``.
-}
decode ::
    -- | The read capability for the application data.
    Cap.Reader ->
    -- | At least as many shares as are required to erasure decode the
    -- ciphertext.
    [(Int, Share)] ->
    -- | An action that results in the ciphertext contained by the shares if
    -- it is possible to recover it, or Nothing.
    IO (Maybe LB.ByteString)
decode :: Reader -> [(Int, Share)] -> IO (Maybe ByteString)
decode Cap.Reader{verifier :: Reader -> Verifier
verifier = Cap.Verifier{Required
required :: Verifier -> Required
required :: Required
required, Required
total :: Verifier -> Required
total :: Required
total, SegmentSize
size :: Verifier -> SegmentSize
size :: SegmentSize
size}} [(Int, Share)]
shares
    | SegmentSize
size SegmentSize -> SegmentSize -> Bool
forall a. Ord a => a -> a -> Bool
> Int64 -> SegmentSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int64 @Integer Int64
forall a. Bounded a => a
maxBound = Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
    | [(Int, Share)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Share)]
shares Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Required -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
required = Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
    | Bool
otherwise = do
        let -- Enough shares to satisfy the ZFEC decoder.
            enoughShares :: [(Int, Share)]
enoughShares = Int -> [(Int, Share)] -> [(Int, Share)]
forall a. Int -> [a] -> [a]
take (Required -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
required) [(Int, Share)]
shares

            -- A list of erasure encoded blocks and positional information.
            -- The outer list gives a share number along with all of the
            -- blocks held in that share.
            blocks :: [(Int, [LB.ByteString])]
            blocks :: [(Int, [ByteString])]
blocks = (Share -> [ByteString]) -> (Int, Share) -> (Int, [ByteString])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Share -> [ByteString]
shareBlocks ((Int, Share) -> (Int, [ByteString]))
-> [(Int, Share)] -> [(Int, [ByteString])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, Share)]
enoughShares

            -- The outer is corresponds to erasure-encoded segments.  The
            -- order corresponds to the order of the segments from the
            -- original input.  Each inner list contains enough blocks to be
            -- erasure-decoded back to a segment.
            explodedBlocks :: [[(Int, LB.ByteString)]]
            explodedBlocks :: [[(Int, ByteString)]]
explodedBlocks = [[(Int, ByteString)]] -> [[(Int, ByteString)]]
forall a. [[a]] -> [[a]]
transpose ([[(Int, ByteString)]] -> [[(Int, ByteString)]])
-> [[(Int, ByteString)]] -> [[(Int, ByteString)]]
forall a b. (a -> b) -> a -> b
$ (Int, [ByteString]) -> [(Int, ByteString)]
fixBlocks ((Int, [ByteString]) -> [(Int, ByteString)])
-> [(Int, [ByteString])] -> [[(Int, ByteString)]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, [ByteString])]
blocks

            -- Figure out how many bytes are expected to be in each segment.
            -- Depending on the ZFEC encoding parameters, it is possible that
            -- we will end up with blocks that are not completely "filled"
            -- with real data.  When these are decoded, we will get _extra_
            -- bytes in the result.  By knowing how many bytes were originally
            -- in our segments, we can recognize and discard these extra
            -- bytes.
            segSize :: Int64
segSize = SegmentSize -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SegmentSize -> Int64)
-> ([(Int, Share)] -> SegmentSize) -> [(Int, Share)] -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parameters -> SegmentSize
paramSegmentSize (Parameters -> SegmentSize)
-> ([(Int, Share)] -> Parameters) -> [(Int, Share)] -> SegmentSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URIExtension -> Parameters
uriExtCodecParams (URIExtension -> Parameters)
-> ([(Int, Share)] -> URIExtension) -> [(Int, Share)] -> Parameters
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Share -> URIExtension
shareURIExtension (Share -> URIExtension)
-> ([(Int, Share)] -> Share) -> [(Int, Share)] -> URIExtension
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Share) -> Share
forall a b. (a, b) -> b
snd ((Int, Share) -> Share)
-> ([(Int, Share)] -> (Int, Share)) -> [(Int, Share)] -> Share
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Share)] -> (Int, Share)
forall a. [a] -> a
head ([(Int, Share)] -> Int64) -> [(Int, Share)] -> Int64
forall a b. (a -> b) -> a -> b
$ [(Int, Share)]
enoughShares

            -- A helper that knows the correct parameters to do ZFEC decoding
            -- for us.
            zunfec' :: [(Int, ByteString)] -> IO ByteString
zunfec' = (Int64 -> ByteString -> ByteString
LB.take Int64
segSize (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (IO ByteString -> IO ByteString)
-> ([(Int, ByteString)] -> IO ByteString)
-> [(Int, ByteString)]
-> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> [(Int, ByteString)] -> IO ByteString
zunfecLazy (Required -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
required) (Required -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Required
total)

        -- Decode every group of blocks back to the original segments.
        [ByteString]
segments <- ([(Int, ByteString)] -> IO ByteString)
-> [[(Int, ByteString)]] -> IO [ByteString]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM [(Int, ByteString)] -> IO ByteString
zunfec' [[(Int, ByteString)]]
explodedBlocks

        -- Combine the segments and perform one more truncation to get the
        -- complete result.  Above where we computed segSize we weren't
        -- careful to find the tail segment size for use with the tail segment
        -- so there might still be some extra bytes in the `segments` list
        -- here.  This additional truncation addresses that.
        Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> ([ByteString] -> ByteString) -> [ByteString] -> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> ByteString -> ByteString
LB.take (SegmentSize -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral SegmentSize
size) (ByteString -> ByteString)
-> ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
LB.concat ([ByteString] -> Maybe ByteString)
-> [ByteString] -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString]
segments
  where
    -- Project the share number out across all of that share's blocks.  The
    -- result is something we can transpose into the correct form for ZFEC
    -- decoding.
    fixBlocks :: (Int, [LB.ByteString]) -> [(Int, LB.ByteString)]
    fixBlocks :: (Int, [ByteString]) -> [(Int, ByteString)]
fixBlocks (Int
sharenum, [ByteString]
bs) = [Int] -> [ByteString] -> [(Int, ByteString)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [Int]
forall a. a -> [a]
repeat Int
sharenum) [ByteString]
bs

makeShareTree :: [MerkleTree] -> MerkleTree
makeShareTree :: [MerkleTree] -> MerkleTree
makeShareTree = [ByteString] -> MerkleTree
makeTreePartial ([ByteString] -> MerkleTree)
-> ([MerkleTree] -> [ByteString]) -> [MerkleTree] -> MerkleTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (MerkleTree -> ByteString) -> [MerkleTree] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map MerkleTree -> ByteString
rootHash

makeCrypttextHash :: Context SHA256 -> CrypttextHash
makeCrypttextHash :: Context SHA256 -> ByteString
makeCrypttextHash = ByteString -> ByteString
sha256 (ByteString -> ByteString)
-> (Context SHA256 -> ByteString) -> Context SHA256 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest SHA256 -> ByteString
toBytes (Digest SHA256 -> ByteString)
-> (Context SHA256 -> Digest SHA256)
-> Context SHA256
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Context SHA256 -> Digest SHA256
forall a. HashAlgorithm a => Context a -> Digest a
hashFinalize
  where
    toBytes :: Digest SHA256 -> ByteString
toBytes = [Word8] -> ByteString
B.pack ([Word8] -> ByteString)
-> (Digest SHA256 -> [Word8]) -> Digest SHA256 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest SHA256 -> [Word8]
forall a. ByteArrayAccess a => a -> [Word8]
BA.unpack

makeCrypttextRootHash :: [CrypttextHash] -> CrypttextHash
makeCrypttextRootHash :: [ByteString] -> ByteString
makeCrypttextRootHash = MerkleTree -> ByteString
rootHash (MerkleTree -> ByteString)
-> ([ByteString] -> MerkleTree) -> [ByteString] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> MerkleTree
makeTreePartial

-- Construct the encoding parameters for the final segment which may be
-- smaller than the earlier segments (if the size of the data to be encoded is
-- not a multiple of the segment size).
-- allmydata.immutable.encode.Encoder._got_all_encoding_parameters
tailParams :: Integral a => Parameters -> a -> Parameters
tailParams :: Parameters -> a -> Parameters
tailParams p :: Parameters
p@Parameters{SegmentSize
paramSegmentSize :: SegmentSize
paramSegmentSize :: Parameters -> SegmentSize
paramSegmentSize, Required
paramRequiredShares :: Required
paramRequiredShares :: Parameters -> Required
paramRequiredShares} a
dataSize =
    Parameters
p{paramSegmentSize :: SegmentSize
paramSegmentSize = Required -> SegmentSize -> SegmentSize
forall m v. (Integral m, Integral v) => m -> v -> v
nextMultipleOf Required
paramRequiredShares SegmentSize
tailSize'}
  where
    tailSize' :: SegmentSize
tailSize' =
        if SegmentSize
tailSize SegmentSize -> SegmentSize -> Bool
forall a. Eq a => a -> a -> Bool
== SegmentSize
0
            then SegmentSize
paramSegmentSize
            else SegmentSize
tailSize
    tailSize :: SegmentSize
tailSize = a -> SegmentSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
dataSize SegmentSize -> SegmentSize -> SegmentSize
forall a. Integral a => a -> a -> a
`mod` SegmentSize
paramSegmentSize

{- | Determine the node numbers of the share tree which are required to verify
 the indicated share number.  The indicated share number is included in the
 result, as are the corresponding hashes from the given tree.
-}
computeNeededShares :: MerkleTree -> Int -> [(Int, B.ByteString)]
computeNeededShares :: MerkleTree -> Int -> [(Int, ByteString)]
computeNeededShares MerkleTree
shareTree Int
sharenum =
    -- In addition to what neededHashes computes we also need to include this
    -- share's own block hash root in the result.  Shove it on the front of
    -- the result here.  This will place it out of order so we'll fix it up
    -- below when we construct the Share.  We also have to translate between
    -- zero-indexed share numbers and 1-indexed leaf numbers.
    --
    -- Is fromJust here safe?  neededHashes returns Nothing when it fails to
    -- compute a merkle proof.  Given the way we're using it, that can
    -- probably only happen if there's a bug inside neededHashes (as opposed
    -- to our passing in some value it doesn't want to provide a result for).
    (MerkleTree -> Int -> Int
leafNumberToNodeNumber MerkleTree
shareTree Int
sharenum Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, MerkleTree -> Int -> ByteString
blockHashRoot MerkleTree
shareTree Int
sharenum) (Int, ByteString) -> [(Int, ByteString)] -> [(Int, ByteString)]
forall a. a -> [a] -> [a]
: Maybe [(Int, ByteString)] -> [(Int, ByteString)]
forall a. HasCallStack => Maybe a -> a
fromJust (MerkleTree -> Int -> Maybe [(Int, ByteString)]
neededHashes MerkleTree
shareTree Int
sharenum)

-- | Find the nth leaf hash in the given tree.
blockHashRoot :: MerkleTree -> Int -> B.ByteString
blockHashRoot :: MerkleTree -> Int -> ByteString
blockHashRoot MerkleTree
tree Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> ByteString
forall a. HasCallStack => String -> a
error String
"Cannot have a negative leaf number"
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
leafs = String -> ByteString
forall a. HasCallStack => String -> a
error String
"Leaf number goes past the end of the tree"
    | Bool
otherwise = [ByteString]
leafs [ByteString] -> Int -> ByteString
forall a. [a] -> Int -> a
!! Int
n
  where
    leafs :: [ByteString]
leafs = MerkleTree -> [ByteString]
leafHashes MerkleTree
tree