{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Crypto.PubKey.MaskGenFunction
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : Good
module Crypto.PubKey.MaskGenFunction (
    MaskGenAlgorithm,
    mgf1,
) where

import Control.Monad (forM_)
import Crypto.Hash (Context, Digest)
import Crypto.Hash.IO (HashAlgorithm (..))
import Crypto.Internal.ByteArray (ByteArray, ByteArrayAccess)
import qualified Crypto.Internal.ByteArray as B
import Data.Bits (shiftR)
import Data.Word (Word8)
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (pokeByteOff)

-- | Represent a mask generation algorithm
type MaskGenAlgorithm seed output =
    seed
    -- ^ seed
    -> Int
    -- ^ length to generate
    -> output

-- | Mask generation algorithm MGF1
mgf1
    :: forall seed output hashAlg
     . (ByteArrayAccess seed, ByteArray output, HashAlgorithm hashAlg)
    => hashAlg
    -> seed
    -> Int
    -> output
mgf1 :: forall seed output hashAlg.
(ByteArrayAccess seed, ByteArray output, HashAlgorithm hashAlg) =>
hashAlg -> seed -> Int -> output
mgf1 hashAlg
hashAlg seed
seed Int
len
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = output
forall a. ByteArray a => a
B.empty
    | Bool
otherwise = Int -> (Ptr Any -> IO ()) -> output
forall a p. ByteArray a => Int -> (Ptr p -> IO ()) -> a
B.allocAndFreeze Int
len ((Ptr Any -> IO ()) -> output) -> (Ptr Any -> IO ()) -> output
forall a b. (a -> b) -> a -> b
$ \Ptr Any
out ->
        -- the seeded context is kept, and copied for each block rather than
        -- rebuilt; each block's digest is written straight into the output
        Int -> (Ptr (Context hashAlg) -> IO ()) -> IO ()
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
ctxSize ((Ptr (Context hashAlg) -> IO ()) -> IO ())
-> (Ptr (Context hashAlg) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr (Context hashAlg)
seeded :: Ptr (Context hashAlg)) ->
            Int -> (Ptr (Context hashAlg) -> IO ()) -> IO ()
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
ctxSize ((Ptr (Context hashAlg) -> IO ()) -> IO ())
-> (Ptr (Context hashAlg) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr (Context hashAlg)
ctx :: Ptr (Context hashAlg)) ->
                Int -> (Ptr (Digest hashAlg) -> IO ()) -> IO ()
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
digestLen ((Ptr (Digest hashAlg) -> IO ()) -> IO ())
-> (Ptr (Digest hashAlg) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr (Digest hashAlg)
dig :: Ptr (Digest hashAlg)) ->
                    Int -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
4 ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr Word8
cnt :: Ptr Word8) -> do
                        Ptr (Context hashAlg) -> IO ()
forall a. HashAlgorithm a => Ptr (Context a) -> IO ()
hashInternalInit Ptr (Context hashAlg)
seeded
                        seed -> (Ptr Word8 -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. seed -> (Ptr p -> IO a) -> IO a
B.withByteArray seed
seed ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p ->
                            Ptr (Context hashAlg) -> Ptr Word8 -> Word32 -> IO ()
forall a.
HashAlgorithm a =>
Ptr (Context a) -> Ptr Word8 -> Word32 -> IO ()
hashInternalUpdate Ptr (Context hashAlg)
seeded Ptr Word8
p (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (seed -> Int
forall ba. ByteArrayAccess ba => ba -> Int
B.length seed
seed))
                        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
maxCounter Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
                            Ptr (Context hashAlg) -> Ptr (Context hashAlg) -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr (Context hashAlg)
ctx Ptr (Context hashAlg)
seeded Int
ctxSize
                            [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
3] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
k ->
                                Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
cnt Int
k (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k))) :: Word8)
                            Ptr (Context hashAlg) -> Ptr Word8 -> Word32 -> IO ()
forall a.
HashAlgorithm a =>
Ptr (Context a) -> Ptr Word8 -> Word32 -> IO ()
hashInternalUpdate Ptr (Context hashAlg)
ctx Ptr Word8
cnt Word32
4
                            Ptr (Context hashAlg) -> Ptr (Digest hashAlg) -> IO ()
forall a.
HashAlgorithm a =>
Ptr (Context a) -> Ptr (Digest a) -> IO ()
hashInternalFinalize Ptr (Context hashAlg)
ctx Ptr (Digest hashAlg)
dig
                            let off :: Int
off = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
digestLen
                            Ptr Any -> Ptr Any -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (Ptr Any
out Ptr Any -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off) (Ptr (Digest hashAlg) -> Ptr Any
forall a b. Ptr a -> Ptr b
castPtr Ptr (Digest hashAlg)
dig) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
digestLen (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off))
  where
    ctxSize :: Int
ctxSize = hashAlg -> Int
forall a. HashAlgorithm a => a -> Int
hashInternalContextSize hashAlg
hashAlg
    digestLen :: Int
digestLen = hashAlg -> Int
forall a. HashAlgorithm a => a -> Int
hashDigestSize hashAlg
hashAlg
    maxCounter :: Int
maxCounter = (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
digestLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
digestLen