mlkem-0.1.0.0: tests/EncapDecap.hs

-- |
-- Module      : EncapDecap
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- Types that are specific to encapsulation/decapsulation test vectors
--
{-# LANGUAGE OverloadedStrings #-}
module EncapDecap
    ( TestGroup(..), Test(..), TestGroupPayload(..)
    , EncapsulationExt(..), DecapsulationExt(..)
    , EncapsulationKeyCheckExt(..), DecapsulationKeyCheckExt(..)
    ) where

import Data.Aeson
import Data.Aeson.Types
import Data.ByteString (ByteString)

import Util

data TestGroup = TestGroup
    { tgId :: Int
    , testType :: String
    , parameterSet :: String
    , function :: String
    , payload :: TestGroupPayload
    } deriving Show

data TestGroupPayload
    = FunctionEncapsulation [Test EncapsulationExt]
    | FunctionDecapsulation [Test DecapsulationExt]
    | FunctionEncapsulationKeyCheck [Test EncapsulationKeyCheckExt]
    | FunctionDecapsulationKeyCheck [Test DecapsulationKeyCheckExt]
    deriving Show

parsePayload :: Object -> Parser TestGroupPayload
parsePayload o = do
    fn <- o .: "function"
    case fn of
        "encapsulation" -> FunctionEncapsulation <$> (o .: "tests")
        "decapsulation" -> FunctionDecapsulation <$> (o .: "tests")
        "encapsulationKeyCheck" -> FunctionEncapsulationKeyCheck <$> (o .: "tests")
        "decapsulationKeyCheck" -> FunctionDecapsulationKeyCheck <$> (o .: "tests")
        unknown -> fail ("parsePayload: unknown function " ++ unknown)

instance FromJSON TestGroup where
    parseJSON = withObject "TestGroup" $ \o -> TestGroup
        <$> o .: "tgId"
        <*> o .: "testType"
        <*> o .: "parameterSet"
        <*> o .: "function"
        <*> parsePayload o

data Test ext = Test
    { tcId :: Int
    , deferred :: Bool
    , tcExt :: ext
    } deriving Show

class TestExt ext where
    parseExt :: Object -> Parser ext

instance TestExt ext => FromJSON (Test ext) where
    parseJSON = withObject "Test" $ \o -> Test
        <$> o .: "tcId"
        <*> o .: "deferred"
        <*> parseExt o

data EncapsulationExt = EncapsulationExt
    { ekEnc :: ByteString
    , dkEnc :: ByteString
    , cEnc :: ByteString
    , kEnc :: ByteString
    , mEnc :: ByteString
    } deriving Show

instance TestExt EncapsulationExt where
    parseExt o = EncapsulationExt
        <$> o .:: "ek"
        <*> o .:: "dk"
        <*> o .:: "c"
        <*> o .:: "k"
        <*> o .:: "m"

data DecapsulationExt = DecapsulationExt
    { ekDec :: ByteString
    , dkDec :: ByteString
    , cDec :: ByteString
    , kDec :: ByteString
    , reasonDec :: String
    } deriving Show

instance TestExt DecapsulationExt where
    parseExt o = DecapsulationExt
        <$> o .:: "ek"
        <*> o .:: "dk"
        <*> o .:: "c"
        <*> o .:: "k"
        <*> o .:  "reason"

data DecapsulationKeyCheckExt = DecapsulationKeyCheckExt
    { passedDkc :: Bool
    , ekDkc :: ByteString
    , dkDkc :: ByteString
    , reasonDkc :: String
    } deriving Show

instance TestExt DecapsulationKeyCheckExt where
    parseExt o = DecapsulationKeyCheckExt
        <$> o .:  "testPassed"
        <*> o .:: "ek"
        <*> o .:: "dk"
        <*> o .:  "reason"

data EncapsulationKeyCheckExt = EncapsulationKeyCheckExt
    { passedEkc :: Bool
    , ekEkc :: ByteString
    , dkEkc :: ByteString
    , reasonEkc :: String
    } deriving Show

instance TestExt EncapsulationKeyCheckExt where
    parseExt o = EncapsulationKeyCheckExt
        <$> o .:  "testPassed"
        <*> o .:: "ek"
        <*> o .:: "dk"
        <*> o .:  "reason"