{-# LANGUAGE OverloadedStrings #-}

module Network.TLS.KeySchedule (
    hkdfExtract,
    hkdfExpandLabel,
    deriveSecret,
) where

import qualified Crypto.Hash as H
import Crypto.KDF.HKDF
import Data.ByteArray (ByteArray, ByteArrayAccess, convert)
import qualified Data.ByteString as BS

import Network.TLS.Crypto
import Network.TLS.Imports
import Network.TLS.Types
import Network.TLS.Wire

----------------------------------------------------------------

-- | @HKDF-Extract@ function.  Returns the pseudorandom key (PRK) from salt and
-- input keying material (IKM).
hkdfExtract
    :: (ByteArray ba, ByteArrayAccess ba) => Hash -> ba -> ba -> ba
hkdfExtract :: forall ba.
(ByteArray ba, ByteArrayAccess ba) =>
Hash -> ba -> ba -> ba
hkdfExtract Hash
SHA1 ba
salt ba
ikm = PRK SHA1 -> ba
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert (ba -> ba -> PRK SHA1
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
extract ba
salt ba
ikm :: PRK H.SHA1)
hkdfExtract Hash
SHA256 ba
salt ba
ikm = PRK SHA256 -> ba
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert (ba -> ba -> PRK SHA256
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
extract ba
salt ba
ikm :: PRK H.SHA256)
hkdfExtract Hash
SHA384 ba
salt ba
ikm = PRK SHA384 -> ba
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert (ba -> ba -> PRK SHA384
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
extract ba
salt ba
ikm :: PRK H.SHA384)
hkdfExtract Hash
SHA512 ba
salt ba
ikm = PRK SHA512 -> ba
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert (ba -> ba -> PRK SHA512
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
extract ba
salt ba
ikm :: PRK H.SHA512)
hkdfExtract Hash
_ ba
_ ba
_ = [Char] -> ba
forall a. HasCallStack => [Char] -> a
error [Char]
"hkdfExtract: unsupported hash"

----------------------------------------------------------------

deriveSecret :: Hash -> Secret -> ByteString -> TranscriptHash -> Secret
deriveSecret :: Hash -> Secret -> ByteString -> TranscriptHash -> Secret
deriveSecret Hash
h Secret
secret ByteString
label (TranscriptHash ByteString
hashedMsgs) =
    Hash -> Secret -> ByteString -> ByteString -> Int -> Secret
forall ba.
(ByteArray ba, ByteArrayAccess ba) =>
Hash -> Secret -> ByteString -> ByteString -> Int -> ba
hkdfExpandLabel Hash
h Secret
secret ByteString
label ByteString
hashedMsgs Int
outlen
  where
    outlen :: Int
outlen = Hash -> Int
hashDigestSize Hash
h

----------------------------------------------------------------

-- | @HKDF-Expand-Label@ function.  Returns output keying material of the
-- specified length from the PRK, customized for a TLS label and context.
hkdfExpandLabel
    :: (ByteArray ba, ByteArrayAccess ba)
    => Hash
    -> Secret
    -> ByteString
    -> ByteString
    -> Int
    -> ba
hkdfExpandLabel :: forall ba.
(ByteArray ba, ByteArrayAccess ba) =>
Hash -> Secret -> ByteString -> ByteString -> Int -> ba
hkdfExpandLabel Hash
h Secret
secret ByteString
label ByteString
ctx Int
outlen = Hash -> Secret -> ByteString -> Int -> ba
forall ba.
(ByteArray ba, ByteArrayAccess ba) =>
Hash -> Secret -> ByteString -> Int -> ba
expand' Hash
h Secret
secret ByteString
hkdfLabel Int
outlen
  where
    hkdfLabel :: ByteString
hkdfLabel = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ do
        Word16 -> Put
putWord16 (Word16 -> Put) -> Word16 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
outlen
        ByteString -> Put
putOpaque8 (ByteString
"tls13 " ByteString -> ByteString -> ByteString
`BS.append` ByteString
label)
        ByteString -> Put
putOpaque8 ByteString
ctx

expand'
    :: (ByteArray ba, ByteArrayAccess ba) => Hash -> Secret -> ByteString -> Int -> ba
expand' :: forall ba.
(ByteArray ba, ByteArrayAccess ba) =>
Hash -> Secret -> ByteString -> Int -> ba
expand' Hash
SHA1 Secret
secret ByteString
label Int
len = PRK SHA1 -> ByteString -> Int -> ba
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
expand (Secret -> PRK SHA1
forall ikm a. ByteArrayAccess ikm => ikm -> PRK a
extractSkip Secret
secret :: PRK H.SHA1) ByteString
label Int
len
expand' Hash
SHA256 Secret
secret ByteString
label Int
len = PRK SHA256 -> ByteString -> Int -> ba
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
expand (Secret -> PRK SHA256
forall ikm a. ByteArrayAccess ikm => ikm -> PRK a
extractSkip Secret
secret :: PRK H.SHA256) ByteString
label Int
len
expand' Hash
SHA384 Secret
secret ByteString
label Int
len = PRK SHA384 -> ByteString -> Int -> ba
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
expand (Secret -> PRK SHA384
forall ikm a. ByteArrayAccess ikm => ikm -> PRK a
extractSkip Secret
secret :: PRK H.SHA384) ByteString
label Int
len
expand' Hash
SHA512 Secret
secret ByteString
label Int
len = PRK SHA512 -> ByteString -> Int -> ba
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
expand (Secret -> PRK SHA512
forall ikm a. ByteArrayAccess ikm => ikm -> PRK a
extractSkip Secret
secret :: PRK H.SHA512) ByteString
label Int
len
expand' Hash
_ Secret
_ ByteString
_ Int
_ = [Char] -> ba
forall a. HasCallStack => [Char] -> a
error [Char]
"expand'"

----------------------------------------------------------------