Support multiple private key formats in generator classes

Signed-off-by: Gilles Peskine <[email protected]>

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: