From 73b5c8ee6ce673f6b7853a94e0dfc37d7b4aafaf Mon Sep 17 00:00:00 2001 From: Gilles Peskine Date: Wed, 8 Jul 2026 12:34:24 +0200 Subject: [PATCH] Support multiple private key formats in generator classes Signed-off-by: Gilles Peskine --- .../mldsa_test_generator.py | 44 +++++++++++-------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/util/mbedtls_maintainer/mldsa_test_generator.py b/util/mbedtls_maintainer/mldsa_test_generator.py index 62dc995cd..615c41b57 100644 --- a/util/mbedtls_maintainer/mldsa_test_generator.py +++ b/util/mbedtls_maintainer/mldsa_test_generator.py @@ -112,10 +112,6 @@ class Generator: def final_arguments(cls) -> List[str]: return [] - @classmethod - def secret_is_seed(cls) -> bool: - return True - @staticmethod def message_for_length(length: int) -> bytes: # b'ABCDE...' wrapping around and repeating every 256 bytes @@ -171,7 +167,7 @@ class Generator: return tc def one_mldsa_sign_deterministic_pure(self, - key: Key, + key: Key, pkf: PrivateKeyFormat, message: bytes, descr: str) -> test_case.TestCase: """Construct one test case for deterministic signature.""" @@ -180,7 +176,7 @@ class Generator: tc.set_function(self.function('sign_message_deterministic', key.kl)) tc.set_dependencies([f'TF_PSA_CRYPTO_PQCP_MLDSA_{key.kl}_ENABLED']) tc.set_arguments(self.metadata_arguments(key.kl, True, True) + [ - test_case.hex_string(key.seed if self.secret_is_seed() else key.secret), + test_case.hex_string(key.private_representation(pkf)), test_case.hex_string(message), test_case.hex_string(signature), ] + self.final_arguments()) @@ -212,12 +208,15 @@ class Generator: def gen_mldsa_pure(self, kl: int) -> Iterator[test_case.TestCase]: """Generate all test cases for pure ML-DSA signature and verification.""" - for i, key in enumerate(KEYS[kl], 1): - yield self.one_mldsa_sign_deterministic_pure(key, MESSAGES[0][0], - f'key#{i}') - for message, descr in MESSAGES[1:]: - yield self.one_mldsa_sign_deterministic_pure(KEYS[kl][0], message, - f'key#1 {descr}') + for pkf, pkf_descr in self.describe_private_key_formats(space_before=True): + for i, key in enumerate(KEYS[kl], 1): + yield self.one_mldsa_sign_deterministic_pure(key, pkf, + MESSAGES[0][0], + f'key#{i}{pkf_descr}') + for message, descr in MESSAGES[1:]: + yield self.one_mldsa_sign_deterministic_pure(KEYS[kl][0], pkf, + message, + f'key#1{pkf_descr} {descr}') for i, key in enumerate(KEYS[kl], 1): yield self.one_mldsa_verify_pure(key, MESSAGES[0][0], True, f'key#{i}') @@ -245,6 +244,9 @@ class Generator: class PQCPGenerator(Generator): """Test mldsa-native entry points.""" + def __init__(self) -> None: + self.private_key_formats = [PrivateKeyFormat.EXPANDED] + @classmethod def function(cls, func: str, kl: int) -> str: if func == 'verify_message': @@ -260,10 +262,6 @@ class PQCPGenerator(Generator): _deterministic: Optional[bool]) -> List[str]: return [] - @classmethod - def secret_is_seed(cls) -> bool: - return False - @staticmethod def one_mldsa_key_pair_from_seed(key: Key, descr: str) -> test_case.TestCase: @@ -329,6 +327,7 @@ class DriverGenerator(Generator): def one_multipart(self, function: str, key: Key, lengths: Sequence[int], + pkf_described: Tuple[Optional[PrivateKeyFormat], str] = (None, ''), tweak_signature: Optional[Callable[[bytes], Tuple[bytes, str, str]]] = None, ) -> test_case.TestCase: """Construct one test case for a multipart operation. @@ -336,7 +335,8 @@ class DriverGenerator(Generator): The number of message chunks must be at most MULTIPART_ARITY. """ message, chunks = self.chunks_for_lengths(lengths, self.MULTIPART_ARITY) - descr = '+'.join(map(str, lengths)) if lengths else 'empty (no update)' + descr = pkf_described[1] + descr += '+'.join(map(str, lengths)) if lengths else 'empty (no update)' type_is_pair = not function.startswith('verif') deterministic = True actual_signature = key.sign_message(message, deterministic=deterministic) @@ -350,8 +350,12 @@ class DriverGenerator(Generator): tc = test_case.TestCase() tc.set_function(self.function(function + '_multipart', key.kl)) tc.set_dependencies([f'TF_PSA_CRYPTO_PQCP_MLDSA_{key.kl}_ENABLED']) + if type_is_pair: + key_bytes = key.private_representation(pkf_described[0]) + else: + key_bytes = key.public tc.set_arguments(self.metadata_arguments(key.kl, type_is_pair, True) + [ - test_case.hex_string(key.seed if type_is_pair else key.public), + test_case.hex_string(key_bytes), str(len(lengths)), ] + [test_case.hex_string(chunk) for chunk in chunks] + [ test_case.hex_string(signature), @@ -379,7 +383,9 @@ class DriverGenerator(Generator): def gen_multipart(self, key: Key) -> Iterator[test_case.TestCase]: """Generate test cases for multipart sign and verify.""" for lengths in self.MANY_MULTIPART_LENGTHS: - yield self.one_multipart('sign_deterministic', key, lengths) + for pkf, pkf_descr in self.describe_private_key_formats(space_after=True): + yield self.one_multipart('sign_deterministic', key, lengths, + pkf_described=(pkf, pkf_descr)) yield self.one_multipart('verify', key, lengths) for descr, func in self.VERIFY_TWEAKS.items(): for lengths in self.FEW_MULTIPART_LENGTHS: