{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | Support the encryption requirements of CHK.
module Tahoe.CHK.Encrypt (encrypt, encryptLazy, decrypt, decryptLazy) where

import Crypto.Cipher.Types (BlockCipher (blockSize, ctrCombine), ivAdd, nullIV)
import Data.ByteArray (ByteArray)
import qualified Data.ByteString.Lazy as LBS
import Data.List (unfoldr)

{- | CTR-mode encrypt a byte string using some block cipher.

 When used for CHKv1 or CHKv2 the block cipher should be AES128.

 This replaces allmydata.immutable.upload.EncryptAnUploadable

 The only noteworthy piece here is that encryption starts with the zero IV.
-}
encrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
encrypt :: cipher -> ba -> ba
encrypt cipher
key = cipher -> IV cipher -> ba -> ba
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> IV cipher -> ba -> ba
ctrCombine cipher
key IV cipher
forall c. BlockCipher c => IV c
nullIV

-- | Like encrypt but operate on lazy bytestrings.
encryptLazy :: forall cipher. BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
encryptLazy :: cipher -> ByteString -> ByteString
encryptLazy cipher
cipher ByteString
lbs = [ByteString] -> ByteString
LBS.concat ([ByteString] -> ByteString)
-> ([ByteString] -> [ByteString]) -> [ByteString] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> ByteString
LBS.fromStrict (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ (IV cipher -> ByteString -> ByteString)
-> [IV cipher] -> [ByteString] -> [ByteString]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (cipher -> IV cipher -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> IV cipher -> ba -> ba
ctrCombine cipher
cipher) [IV cipher]
ivs [ByteString]
blocks
  where
    -- The underlying encryption function works on strict bytes.  Here's the
    -- number of *blocks* to feed to it (that is, to make strict) at a time.
    -- This value here is a magic number that is meant to represent a good
    -- compromise between performance and number of bytes forced at one time.
    workingBlocks :: Int
workingBlocks = Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
16

    -- The size of a block is determined by the cipher.
    workingBytes :: Int
workingBytes = Int
workingBlocks Int -> Int -> Int
forall a. Num a => a -> a -> a
* cipher -> Int
forall cipher. BlockCipher cipher => cipher -> Int
blockSize @cipher cipher
forall a. HasCallStack => a
undefined

    ivs :: [IV cipher]
ivs = (IV cipher -> IV cipher) -> IV cipher -> [IV cipher]
forall a. (a -> a) -> a -> [a]
iterate (IV cipher -> Int -> IV cipher
forall c. IV c -> Int -> IV c
`ivAdd` Int
workingBlocks) IV cipher
forall c. BlockCipher c => IV c
nullIV
    blocks :: [ByteString]
blocks = ByteString -> ByteString
LBS.toStrict (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ByteString -> Maybe (ByteString, ByteString))
-> ByteString -> [ByteString]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ByteString -> Maybe (ByteString, ByteString)
takeChunk ByteString
lbs

    takeChunk :: ByteString -> Maybe (ByteString, ByteString)
takeChunk ByteString
"" = Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing
    takeChunk ByteString
xs = (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just ((ByteString, ByteString) -> Maybe (ByteString, ByteString))
-> (ByteString -> (ByteString, ByteString))
-> ByteString
-> Maybe (ByteString, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> ByteString -> (ByteString, ByteString)
LBS.splitAt (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
workingBytes) (ByteString -> Maybe (ByteString, ByteString))
-> ByteString -> Maybe (ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString
xs

-- | AES128-CTR decrypt a byte string in the manner used by CHK.
decrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
decrypt :: cipher -> ba -> ba
decrypt = cipher -> ba -> ba
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
encrypt

-- | Like decrypt but operate on lazy bytestrings.
decryptLazy :: BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
decryptLazy :: cipher -> ByteString -> ByteString
decryptLazy = cipher -> ByteString -> ByteString
forall cipher.
BlockCipher cipher =>
cipher -> ByteString -> ByteString
encryptLazy