HKDF.hs

{-# LANGUAGE BangPatterns #-}

-- |
-- Module      : Crypto.KDF.HKDF
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : unknown
--
-- Key Derivation Function based on HMAC
--
-- See RFC5869
module Crypto.KDF.HKDF (
    PRK,
    extract,
    extractSkip,
    expand,
    toPRK,
) where

import Crypto.Hash
import Crypto.Internal.ByteArray (
    ByteArray,
    ByteArrayAccess,
    ScrubbedBytes,
 )
import qualified Crypto.Internal.ByteArray as B
import Crypto.MAC.HMAC
import Data.Word

-- | Pseudo Random Key
data PRK a = PRK (HMAC a) | PRK_NoExpand ScrubbedBytes
    deriving (Eq)

instance ByteArrayAccess (PRK a) where
    length (PRK hm) = B.length hm
    length (PRK_NoExpand sb) = B.length sb
    withByteArray (PRK hm) = B.withByteArray hm
    withByteArray (PRK_NoExpand sb) = B.withByteArray sb

-- | Extract a Pseudo Random Key using the parameter and the underlaying hash mechanism
extract
    :: (HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm)
    => salt
    -- ^ Salt
    -> ikm
    -- ^ Input Keying Material
    -> PRK a
    -- ^ Pseudo random key
extract salt ikm = PRK $ hmac salt ikm

-- | Create a PRK directly from the input key material.
--
-- Only use when guaranteed to have a good quality and random data to use directly as key.
-- This effectively skip a HMAC with key=salt and data=key.
extractSkip
    :: ByteArrayAccess ikm
    => ikm
    -> PRK a
extractSkip ikm = PRK_NoExpand $ B.convert ikm

-- | Expand key material of specific length out of the parameters
expand
    :: (HashAlgorithm a, ByteArrayAccess info, ByteArray out)
    => PRK a
    -- ^ Pseudo Random Key
    -> info
    -- ^ Optional context and application specific information
    -> Int
    -- ^ Output length in bytes
    -> out
    -- ^ Output data
expand prkAt infoAt outputLength =
    let hF = hFGet prkAt
     in B.concat $ loop hF B.empty outputLength 1
  where
    hFGet :: (HashAlgorithm a, ByteArrayAccess b) => PRK a -> (b -> HMAC a)
    hFGet prk = case prk of
        PRK hmacKey -> hmac hmacKey
        PRK_NoExpand ikm -> hmac ikm

    info :: ScrubbedBytes
    info = B.convert infoAt

    loop
        :: HashAlgorithm a
        => (ScrubbedBytes -> HMAC a)
        -> ScrubbedBytes
        -> Int
        -> Word8
        -> [ScrubbedBytes]
    loop hF tim1 n i
        | n <= 0 = []
        | otherwise =
            let input = B.concat [tim1, info, B.singleton i] :: ScrubbedBytes
                ti = B.convert $ hF input
                hashLen = B.length ti
                r = n - hashLen
             in (if n >= hashLen then ti else B.take n ti)
                    : loop hF ti r (i + 1)

toPRK :: (HashAlgorithm a, ByteArrayAccess ba) => ba -> Maybe (PRK a)
toPRK bs = case digestFromByteString bs of
    Nothing -> Nothing
    Just digest -> Just $ PRK $ HMAC digest