{-# LANGUAGE ScopedTypeVariables #-}

{- | Implement the scheme for encoding ciphertext into SDMF shares (and
 decoding it again).
-}
module Tahoe.SDMF.Internal.Encoding where

import Control.Monad (when)
import Control.Monad.Fail (MonadFail)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Crypto.Hash (digestFromByteString)
import Crypto.Random (MonadRandom)
import Data.Bifunctor (Bifunctor (bimap))
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Int (Int64)
import qualified Data.Text as T
import Data.Word (Word16, Word64, Word8)
import Tahoe.CHK (padCiphertext, zfec, zunfec)
import Tahoe.CHK.Merkle (MerkleTree (MerkleLeaf))
import Tahoe.CHK.SHA256d (Digest' (Digest'), zero)
import Tahoe.SDMF.Internal.Capability (Reader (..), Writer (..), deriveReader)
import Tahoe.SDMF.Internal.Converting (from, tryInto)
import qualified Tahoe.SDMF.Internal.Keys as Keys
import Tahoe.SDMF.Internal.Share (HashChain (HashChain), Share (..))

{- | Given a pre-determined key pair and sequence number, encode some
 ciphertext into a collection of SDMF shares.

 A key pair *uniquely identifies* a "slot" (the storage location for the shares).
 Thus they cannot be re-used for "different" data.  Any shares created with a
 given key pair are part of the same logical data object.
-}
encode :: (MonadFail m, MonadIO m, MonadRandom m) => Keys.KeyPair -> Keys.SDMF_IV -> Word64 -> Word16 -> Word16 -> LB.ByteString -> m ([Share], Writer)
encode :: KeyPair
-> SDMF_IV
-> Word64
-> Word16
-> Word16
-> ByteString
-> m ([Share], Writer)
encode KeyPair
keypair SDMF_IV
iv Word64
shareSequenceNumber Word16
required Word16
total ByteString
ciphertext = do
    -- Make sure the encoding parameters fit into a Word8
    Word8
requiredAsWord8 <- String -> Word16 -> m Word8
forall b a (m :: * -> *). TryFrom a b m => String -> a -> m b
tryInto @Word8 (String
"must have 0 < required < 255 but required == " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word16 -> String
forall a. Show a => a -> String
show Word16
required) Word16
required
    Word8
totalAsWord8 <- String -> Word16 -> m Word8
forall b a (m :: * -> *). TryFrom a b m => String -> a -> m b
tryInto @Word8 (String
"must have 0 < total < 256 but total == " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word16 -> String
forall a. Show a => a -> String
show Word16
total) Word16
total

    -- And that they make sense together.
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word16
required Word16 -> Word16 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word16
total) (String -> m ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"must have required < total but required == " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word16 -> String
forall a. Show a => a -> String
show Word16
required String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", total == " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word16 -> String
forall a. Show a => a -> String
show Word16
total)

    -- They look okay, we can proceed.
    [ByteString]
blocks <- IO [ByteString] -> m [ByteString]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [ByteString] -> m [ByteString])
-> IO [ByteString] -> m [ByteString]
forall a b. (a -> b) -> a -> b
$ (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
LB.fromStrict ([ByteString] -> [ByteString])
-> IO [ByteString] -> IO [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> ByteString -> IO [ByteString]
zfec (Word16 -> Int
forall a b. From a b => a -> b
from Word16
required) (Word16 -> Int
forall a b. From a b => a -> b
from Word16
total) ByteString
paddedCiphertext

    -- We know the length won't be negative (doesn't make sense) and we
    -- know all positive values fit into a Word64 so we can do this
    -- conversion safely.  But if it needs to fail for some reason, it
    -- can do so safely.
    Word64
dataLength <- String -> Int64 -> m Word64
forall b a (m :: * -> *). TryFrom a b m => String -> a -> m b
tryInto @Word64 String
"must have 0 <= data length" (ByteString -> Int64
LB.length ByteString
ciphertext)

    -- All segments are the same so we can figure the size from any one
    -- block.  This conversion might fail because of Int64 vs Word64 but
    -- only for truly, truly tremendous share data.
    Word64
shareSegmentSize <- String -> Int64 -> m Word64
forall b a (m :: * -> *). TryFrom a b m => String -> a -> m b
tryInto @Word64 String
"must have segment size < 2^63" (ByteString -> Int64
LB.length ([ByteString] -> ByteString
forall a. [a] -> a
head [ByteString]
blocks))

    let makeShare' :: ByteString -> ByteString -> Share
makeShare' =
            (ByteString -> ByteString -> Share)
-> ByteString -> ByteString -> Share
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((ByteString -> ByteString -> Share)
 -> ByteString -> ByteString -> Share)
-> (ByteString -> ByteString -> Share)
-> ByteString
-> ByteString
-> Share
forall a b. (a -> b) -> a -> b
$
                Word64
-> SDMF_IV
-> Word8
-> Word8
-> Word64
-> Word64
-> Verification
-> ByteString
-> ByteString
-> Share
makeShare
                    Word64
shareSequenceNumber
                    SDMF_IV
iv
                    Word8
requiredAsWord8
                    Word8
totalAsWord8
                    Word64
dataLength
                    Word64
shareSegmentSize
                    (KeyPair -> Verification
Keys.toVerificationKey KeyPair
keypair)

    let makeShare'' :: [ByteString -> Share]
makeShare'' = ByteString -> ByteString -> Share
makeShare' (ByteString -> ByteString -> Share)
-> [ByteString] -> [ByteString -> Share]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ByteString]
blocks

        resultE :: Either T.Text [Share]
        resultE :: Either Text [Share]
resultE = (((ByteString -> Share) -> Either Text Share)
-> [ByteString -> Share] -> Either Text [Share]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (((ByteString -> Share) -> Either Text Share)
 -> [ByteString -> Share] -> Either Text [Share])
-> (Either Text ByteString
    -> (ByteString -> Share) -> Either Text Share)
-> Either Text ByteString
-> [ByteString -> Share]
-> Either Text [Share]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ByteString -> Share)
 -> Either Text ByteString -> Either Text Share)
-> Either Text ByteString
-> (ByteString -> Share)
-> Either Text Share
forall a b c. (a -> b -> c) -> b -> a -> c
flip (ByteString -> Share)
-> Either Text ByteString -> Either Text Share
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) Either Text ByteString
encryptedPrivateKey [ByteString -> Share]
makeShare''
    (Text -> m ([Share], Writer))
-> (([Share], Writer) -> m ([Share], Writer))
-> Either Text ([Share], Writer)
-> m ([Share], Writer)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> m ([Share], Writer)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m ([Share], Writer))
-> (Text -> String) -> Text -> m ([Share], Writer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) ([Share], Writer) -> m ([Share], Writer)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((,) ([Share] -> Writer -> ([Share], Writer))
-> Either Text [Share] -> Either Text (Writer -> ([Share], Writer))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text [Share]
resultE Either Text (Writer -> ([Share], Writer))
-> Either Text Writer -> Either Text ([Share], Writer)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Either Text Writer
cap)
  where
    paddedCiphertext :: ByteString
paddedCiphertext = ByteString -> ByteString
LB.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Word16 -> ByteString -> ByteString
padCiphertext Word16
required ByteString
ciphertext
    -- We can compute a capability immediately.
    cap :: Either Text Writer
cap = KeyPair -> Either Text Writer
capabilityForKeyPair KeyPair
keypair
    encryptedPrivateKey :: Either Text ByteString
encryptedPrivateKey = (Write -> Signature -> ByteString)
-> Signature -> Write -> ByteString
forall a b c. (a -> b -> c) -> b -> a -> c
flip Write -> Signature -> ByteString
Keys.encryptSignatureKey (KeyPair -> Signature
Keys.toSignatureKey KeyPair
keypair) (Write -> ByteString)
-> Either Text Write -> Either Text ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Writer -> Write
writerWriteKey (Writer -> Write) -> Either Text Writer -> Either Text Write
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text Writer
cap)

makeShare ::
    Word64 ->
    Keys.SDMF_IV ->
    Word8 ->
    Word8 ->
    Word64 ->
    Word64 ->
    Keys.Verification ->
    B.ByteString ->
    LB.ByteString ->
    Share
makeShare :: Word64
-> SDMF_IV
-> Word8
-> Word8
-> Word64
-> Word64
-> Verification
-> ByteString
-> ByteString
-> Share
makeShare Word64
shareSequenceNumber SDMF_IV
shareIV Word8
shareRequiredShares Word8
shareTotalShares Word64
shareDataLength Word64
shareSegmentSize Verification
shareVerificationKey ByteString
shareEncryptedPrivateKey ByteString
shareData = Share :: Word64
-> ByteString
-> SDMF_IV
-> Word8
-> Word8
-> Word64
-> Word64
-> Verification
-> ByteString
-> HashChain
-> MerkleTree ByteString SHA256d
-> ByteString
-> ByteString
-> Share
Share{Word8
Word64
ByteString
ByteString
MerkleTree ByteString SHA256d
SDMF_IV
Verification
HashChain
forall value. MerkleTree value SHA256d
shareEncryptedPrivateKey :: ByteString
shareData :: ByteString
shareBlockHashTree :: MerkleTree ByteString SHA256d
shareHashChain :: HashChain
shareSignature :: ByteString
shareVerificationKey :: Verification
shareDataLength :: Word64
shareSegmentSize :: Word64
shareRequiredShares :: Word8
shareTotalShares :: Word8
shareIV :: SDMF_IV
shareRootHash :: ByteString
shareSequenceNumber :: Word64
shareBlockHashTree :: forall value. MerkleTree value SHA256d
shareHashChain :: HashChain
shareSignature :: ByteString
shareRootHash :: ByteString
shareData :: ByteString
shareEncryptedPrivateKey :: ByteString
shareVerificationKey :: Verification
shareSegmentSize :: Word64
shareDataLength :: Word64
shareTotalShares :: Word8
shareRequiredShares :: Word8
shareIV :: SDMF_IV
shareSequenceNumber :: Word64
..}
  where
    shareRootHash :: ByteString
shareRootHash = Int -> Word8 -> ByteString
B.replicate Int
32 Word8
0
    shareSignature :: ByteString
shareSignature = Int -> Word8 -> ByteString
B.replicate Int
32 Word8
0 -- XXX Actually compute sig, and is it 32 bytes?
    shareHashChain :: HashChain
shareHashChain = [(Word16, ByteString)] -> HashChain
HashChain []
    shareBlockHashTree :: MerkleTree value SHA256d
shareBlockHashTree = Digest' SHA256d -> MerkleTree value SHA256d
forall value hash. Digest' hash -> MerkleTree value hash
MerkleLeaf Digest' SHA256d
forall hash. HashAlgorithm hash => Digest' hash
zero -- XXX Real hash here, plus length check

{- | Decode some SDMF shares to recover the original ciphertext.

 TODO: Use the read capability to verify the shares were constructed with
 information from the matching write capability.
-}
decode :: (MonadFail m, MonadIO m) => Reader -> [(Word16, Share)] -> m LB.ByteString
decode :: Reader -> [(Word16, Share)] -> m ByteString
decode Reader
_ [] = String -> m ByteString
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Cannot decode with no shares"
decode Reader
_ s :: [(Word16, Share)]
s@((Word16
_, Share{Word8
shareRequiredShares :: Word8
shareRequiredShares :: Share -> Word8
shareRequiredShares, Word8
shareTotalShares :: Word8
shareTotalShares :: Share -> Word8
shareTotalShares, Word64
shareDataLength :: Word64
shareDataLength :: Share -> Word64
shareDataLength}) : [(Word16, Share)]
shares)
    -- Make sure we have enough shares.
    | [(Word16, Share)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Word16, Share)]
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
requiredAsInt =
        String -> m ByteString
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m ByteString) -> String -> m ByteString
forall a b. (a -> b) -> a -> b
$ String
"got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([(Word16, Share)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Word16, Share)]
shares) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" shares, required " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String
forall a. Show a => a -> String
show Word8
shareRequiredShares
    | Bool
otherwise = do
        -- Make sure this implementation can handle the amount of data involved.
        -- Since we use lazy ByteString we're limited to 2^63-1 bytes rather than
        -- 2^64-1 bytes so there are some SDMF shares we can't interpret right
        -- now.
        Int64
shareDataLength' <- String -> Word64 -> m Int64
forall b a (m :: * -> *). TryFrom a b m => String -> a -> m b
tryInto @Int64 (String
"share data length " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word64 -> String
forall a. Show a => a -> String
show Word64
shareDataLength String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" is beyond maximum supported by this implementation " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int64 -> String
forall a. Show a => a -> String
show (Int64
forall a. Bounded a => a
maxBound :: Int64)) Word64
shareDataLength
        ByteString
ciphertext <- IO ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [(Int, ByteString)] -> IO ByteString
zunfec Int
requiredAsInt Int
totalAsInt (Int -> [(Int, ByteString)] -> [(Int, ByteString)]
forall a. Int -> [a] -> [a]
take Int
requiredAsInt [(Int, ByteString)]
blocks)
        ByteString -> m ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> m ByteString)
-> (ByteString -> ByteString) -> ByteString -> m ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> ByteString -> ByteString
LB.take Int64
shareDataLength' (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LB.fromStrict (ByteString -> m ByteString) -> ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
ciphertext
  where
    blocks :: [(Int, ByteString)]
blocks = (Word16 -> Int)
-> (Share -> ByteString) -> (Word16, Share) -> (Int, ByteString)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (forall b. From Word16 b => Word16 -> b
forall a b. From a b => a -> b
from @Word16) (ByteString -> ByteString
LB.toStrict (ByteString -> ByteString)
-> (Share -> ByteString) -> Share -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Share -> ByteString
shareData) ((Word16, Share) -> (Int, ByteString))
-> [(Word16, Share)] -> [(Int, ByteString)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Word16, Share)]
s

    requiredAsInt :: Int
requiredAsInt = Word8 -> Int
forall a b. From a b => a -> b
from Word8
shareRequiredShares
    totalAsInt :: Int
totalAsInt = Word8 -> Int
forall a b. From a b => a -> b
from Word8
shareTotalShares

-- | Compute an SDMF write capability for a given keypair.
capabilityForKeyPair :: Keys.KeyPair -> Either T.Text Writer
capabilityForKeyPair :: KeyPair -> Either Text Writer
capabilityForKeyPair KeyPair
keypair =
    Write -> Reader -> Writer
Writer (Write -> Reader -> Writer)
-> Either Text Write -> Either Text (Reader -> Writer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text Write
writerWriteKey Either Text (Reader -> Writer)
-> Either Text Reader -> Either Text Writer
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Either Text (Maybe Reader) -> Either Text Reader
forall e a. e -> Either e (Maybe a) -> Either e a
maybeToEither' Text
"Failed to derive read capability" Either Text (Maybe Reader)
writerReader
  where
    writerWriteKey :: Either Text Write
writerWriteKey = Text -> Maybe Write -> Either Text Write
forall a b. a -> Maybe b -> Either a b
maybeToEither Text
"Failed to derive write key" (Maybe Write -> Either Text Write)
-> (KeyPair -> Maybe Write) -> KeyPair -> Either Text Write
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Signature -> Maybe Write
Keys.deriveWriteKey (Signature -> Maybe Write)
-> (KeyPair -> Signature) -> KeyPair -> Maybe Write
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KeyPair -> Signature
Keys.toSignatureKey (KeyPair -> Either Text Write) -> KeyPair -> Either Text Write
forall a b. (a -> b) -> a -> b
$ KeyPair
keypair
    verificationKeyHash :: Maybe (Digest' SHA256d)
verificationKeyHash = (Digest SHA256d -> Digest' SHA256d)
-> Maybe (Digest SHA256d) -> Maybe (Digest' SHA256d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Digest SHA256d -> Digest' SHA256d
forall a. Digest a -> Digest' a
Digest' (Maybe (Digest SHA256d) -> Maybe (Digest' SHA256d))
-> (KeyPair -> Maybe (Digest SHA256d))
-> KeyPair
-> Maybe (Digest' SHA256d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe (Digest SHA256d)
forall a ba.
(HashAlgorithm a, ByteArrayAccess ba) =>
ba -> Maybe (Digest a)
digestFromByteString (ByteString -> Maybe (Digest SHA256d))
-> (KeyPair -> ByteString) -> KeyPair -> Maybe (Digest SHA256d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Verification -> ByteString
Keys.deriveVerificationHash (Verification -> ByteString)
-> (KeyPair -> Verification) -> KeyPair -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KeyPair -> Verification
Keys.toVerificationKey (KeyPair -> Maybe (Digest' SHA256d))
-> KeyPair -> Maybe (Digest' SHA256d)
forall a b. (a -> b) -> a -> b
$ KeyPair
keypair
    writerReader :: Either Text (Maybe Reader)
writerReader = Write -> Digest' SHA256d -> Maybe Reader
deriveReader (Write -> Digest' SHA256d -> Maybe Reader)
-> Either Text Write
-> Either Text (Digest' SHA256d -> Maybe Reader)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either Text Write
writerWriteKey Either Text (Digest' SHA256d -> Maybe Reader)
-> Either Text (Digest' SHA256d) -> Either Text (Maybe Reader)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Maybe (Digest' SHA256d) -> Either Text (Digest' SHA256d)
forall a b. a -> Maybe b -> Either a b
maybeToEither Text
"Failed to interpret verification hash" Maybe (Digest' SHA256d)
verificationKeyHash

maybeToEither :: a -> Maybe b -> Either a b
maybeToEither :: a -> Maybe b -> Either a b
maybeToEither a
a Maybe b
Nothing = a -> Either a b
forall a b. a -> Either a b
Left a
a
maybeToEither a
_ (Just b
b) = b -> Either a b
forall a b. b -> Either a b
Right b
b

maybeToEither' :: e -> Either e (Maybe a) -> Either e a
maybeToEither' :: e -> Either e (Maybe a) -> Either e a
maybeToEither' e
e (Right Maybe a
Nothing) = e -> Either e a
forall a b. a -> Either a b
Left e
e
maybeToEither' e
_ (Right (Just a
r)) = a -> Either e a
forall a b. b -> Either a b
Right a
r
maybeToEither' e
_ (Left e
e) = e -> Either e a
forall a b. a -> Either a b
Left e
e