Skip to content

ML-DSA test generator: support expanded and joined key formats #312

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gilles-peskine-arm wants to merge 3 commits into
base: main
Choose a base branch
from
Open

ML-DSA test generator: support expanded and joined key formats #312

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3a97573
Ground support for multiple key formats
gilles-peskine-arm Jun 22, 2026
73b5c8e
Support multiple private key formats in generator classes
gilles-peskine-arm Jul 8, 2026
98fc1f1
Support multiple private key formats in DriverGenerator
gilles-peskine-arm Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 74 additions & 20 deletions util/mbedtls_maintainer/mldsa_test_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
# SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later

import collections
import enum
import functools
from typing import Callable, Iterator, List, Optional, Sequence, Tuple

Expand Down Expand Up @@ -34,6 +35,22 @@
b'\x00' * 32,
]


class PrivateKeyFormat(enum.Enum):
"""Key representation to pass to the signature function."""
SEED = 0 # key = 32-byte seed
EXPANDED = 1 # key in the standard expanded format
SEED_PLUS_EXPANDED = 2 # key = join(seed, standard-expanded-format)

def descr(self) -> str:
"""A short textual description of the format."""
return [
'seed',
'expanded',
'joined',
][self.value]


class Key:
"""An MLDSA key pair."""
#pylint: disable=too-few-public-methods
Expand All @@ -43,6 +60,17 @@ def __init__(self, kl: int, seed: bytes) -> None:
self.seed = seed
self.public, self.secret = PURE[kl]._keygen_internal(seed)

def private_representation(self,
pkf: Optional[PrivateKeyFormat] = None) -> bytes:
if pkf == PrivateKeyFormat.SEED or pkf is None:
return self.seed
elif pkf == PrivateKeyFormat.EXPANDED:
return self.secret
elif pkf == PrivateKeyFormat.SEED_PLUS_EXPANDED:
return self.seed + self.secret
else:
raise ValueError(f'Unsupported private key format: {pkf}')

@functools.lru_cache(maxsize=9999)
def sign_message(self, message: bytes, deterministic: bool) -> bytes:
PURE[self.kl].set_drbg_seed(bytes(48))
Expand All @@ -66,6 +94,9 @@ def sign_message(self, message: bytes, deterministic: bool) -> bytes:
class Generator:
"""Abstract base class to generate tests for one API."""

def __init__(self) -> None:
self.private_key_formats: Sequence[PrivateKeyFormat] = [PrivateKeyFormat.SEED]

@classmethod
def function(cls, func: str, kl: int) -> str:
raise NotImplementedError
Expand All @@ -81,15 +112,23 @@ def metadata_arguments(cls,
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
return bytes(b & 0xff for b in range(65, 65 + length))

def describe_private_key_formats(self,
space_before: bool = False,
space_after: bool = False,
) -> Iterator[Tuple[PrivateKeyFormat, str]]:
if len(self.private_key_formats) == 1:
yield (self.private_key_formats[0], '')
return
for pkf in self.private_key_formats:
yield (pkf, ''.join([' ' if space_before else '',
'(', pkf.descr(), ')',
' ' if space_after else '']))

def chunks_for_lengths(self,
lengths: Sequence[int],
arity: Optional[int] = None,
Expand Down Expand Up @@ -128,7 +167,7 @@ def one_mldsa_public_key_from_seed(self, key: Key,
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."""
Expand All @@ -137,7 +176,7 @@ def one_mldsa_sign_deterministic_pure(self,
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())
Expand Down Expand Up @@ -169,12 +208,15 @@ def one_mldsa_verify_pure(self,

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}')
Expand Down Expand Up @@ -202,6 +244,9 @@ def gen_all(self) -> Iterator[test_case.TestCase]:
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':
Expand All @@ -217,10 +262,6 @@ def metadata_arguments(cls,
_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:
Expand Down Expand Up @@ -286,14 +327,16 @@ def gen_key_management(self, kl: int) -> Iterator[test_case.TestCase]:

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.

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)
Expand All @@ -307,8 +350,12 @@ def one_multipart(self, function: str, key: Key,
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),
Expand Down Expand Up @@ -336,7 +383,9 @@ def one_multipart(self, function: str, key: Key,
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:
Expand All @@ -345,8 +394,13 @@ def gen_multipart(self, key: Key) -> Iterator[test_case.TestCase]:
yield self.one_multipart('verify', key, lengths,
tweak_signature=tweak)

def gen_all(self, multipart: bool = False) -> Iterator[test_case.TestCase]:
def gen_all(self,
multipart: bool = False,
private_key_formats: Optional[Sequence[PrivateKeyFormat]] = None,
) -> Iterator[test_case.TestCase]:
"""Generate all the tests for this API."""
if private_key_formats is not None:
self.private_key_formats = private_key_formats
yield from super().gen_all()
if multipart:
for kl in sorted(KEYS.keys()):
Expand Down