{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Crypto.HPKE.KeySchedule (
    -- * Types
    Mode (..),

    -- * Key schedule
    keySchedule,
) where

import Crypto.KDF.HKDF (toPRK)
import qualified Data.ByteString as BS

import Crypto.HPKE.KDF
import Crypto.HPKE.Types

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

data Mode
    = ModeBase
    | ModePsk
    | ModeAuth
    | ModeAuthPsk
    deriving (Mode -> Mode -> Bool
(Mode -> Mode -> Bool) -> (Mode -> Mode -> Bool) -> Eq Mode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Mode -> Mode -> Bool
== :: Mode -> Mode -> Bool
$c/= :: Mode -> Mode -> Bool
/= :: Mode -> Mode -> Bool
Eq, Int -> Mode -> ShowS
[Mode] -> ShowS
Mode -> String
(Int -> Mode -> ShowS)
-> (Mode -> String) -> ([Mode] -> ShowS) -> Show Mode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Mode -> ShowS
showsPrec :: Int -> Mode -> ShowS
$cshow :: Mode -> String
show :: Mode -> String
$cshowList :: [Mode] -> ShowS
showList :: [Mode] -> ShowS
Show)

{- FOURMOLU_DISABLE -}
fromMode :: Mode -> Word8
fromMode :: Mode -> Word8
fromMode Mode
ModeBase    = Word8
0x00
fromMode Mode
ModePsk     = Word8
0x01
fromMode Mode
ModeAuth    = Word8
0x02
fromMode Mode
ModeAuthPsk = Word8
0x03
{- FOURMOLU_ENABLE -}

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

keySchedule
    :: forall h
     . (HashAlgorithm h, KDF h)
    => h
    -> Suite
    -> Int
    -> Int
    -> Mode
    -> Info
    -> PSK
    -> PSK_ID
    -> SharedSecret
    -> Either HPKEError (Key, Nonce, Int, PRK h)
keySchedule :: forall h.
(HashAlgorithm h, KDF h) =>
h
-> Suite
-> Int
-> Int
-> Mode
-> Suite
-> Suite
-> Suite
-> SharedSecret
-> Either HPKEError (Suite, Suite, Int, PRK h)
keySchedule h
h Suite
suite Int
nk Int
nn Mode
mode Suite
info Suite
psk Suite
psk_id SharedSecret
shared_secret =
    case Suite -> Maybe (PRK h)
forall a ba.
(HashAlgorithm a, ByteArrayAccess ba) =>
ba -> Maybe (PRK a)
toPRK Suite
exporter_secret of
        Maybe (PRK h)
Nothing -> HPKEError -> Either HPKEError (Suite, Suite, Int, PRK h)
forall a b. a -> Either a b
Left (HPKEError -> Either HPKEError (Suite, Suite, Int, PRK h))
-> HPKEError -> Either HPKEError (Suite, Suite, Int, PRK h)
forall a b. (a -> b) -> a -> b
$ String -> HPKEError
KeyScheduleError String
"cannot convert to PRK"
        Just PRK h
prk -> (Suite, Suite, Int, PRK h)
-> Either HPKEError (Suite, Suite, Int, PRK h)
forall a b. b -> Either a b
Right (Suite
key, Suite
base_nonce, Int
0, PRK h
prk)
  where
    psk_id_hash :: PRK h
psk_id_hash = Suite -> Suite -> Suite -> Suite -> PRK h
forall h. KDF h => Suite -> Suite -> Suite -> Suite -> PRK h
labeledExtract Suite
suite Suite
"" Suite
"psk_id_hash" Suite
psk_id :: PRK h
    info_hash :: PRK h
info_hash = Suite -> Suite -> Suite -> Suite -> PRK h
forall h. KDF h => Suite -> Suite -> Suite -> Suite -> PRK h
labeledExtract Suite
suite Suite
"" Suite
"info_hash" Suite
info :: PRK h
    key_schedule_context :: Suite
key_schedule_context =
        Word8 -> Suite
BS.singleton (Mode -> Word8
fromMode Mode
mode) Suite -> Suite -> Suite
forall a. Semigroup a => a -> a -> a
<> PRK h -> Suite
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert PRK h
psk_id_hash Suite -> Suite -> Suite
forall a. Semigroup a => a -> a -> a
<> PRK h -> Suite
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert PRK h
info_hash
            :: ByteString

    secret :: PRK h
secret = Suite -> Suite -> Suite -> Suite -> PRK h
forall h. KDF h => Suite -> Suite -> Suite -> Suite -> PRK h
labeledExtract Suite
suite (SharedSecret -> Suite
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
convert SharedSecret
shared_secret) Suite
"secret" Suite
psk :: PRK h

    key :: Suite
key = Suite -> PRK h -> Suite -> Suite -> Int -> Suite
forall h. KDF h => Suite -> PRK h -> Suite -> Suite -> Int -> Suite
labeledExpand Suite
suite PRK h
secret Suite
"key" Suite
key_schedule_context Int
nk
    base_nonce :: Suite
base_nonce = Suite -> PRK h -> Suite -> Suite -> Int -> Suite
forall h. KDF h => Suite -> PRK h -> Suite -> Suite -> Int -> Suite
labeledExpand Suite
suite PRK h
secret Suite
"base_nonce" Suite
key_schedule_context Int
nn

    exporter_secret :: Suite
exporter_secret = Suite -> PRK h -> Suite -> Suite -> Int -> Suite
forall h. KDF h => Suite -> PRK h -> Suite -> Suite -> Int -> Suite
labeledExpand Suite
suite PRK h
secret Suite
"exp" Suite
key_schedule_context (Int -> Suite) -> Int -> Suite
forall a b. (a -> b) -> a -> b
$ h -> Int
forall a. HashAlgorithm a => a -> Int
hashDigestSize h
h