Add other versions of dilithium to the key manager.

PiperOrigin-RevId: 392376467

7 files changed, 233 insertions, 67 deletions

diff --git a/cc/experimental/pqcrypto/signature/dilithium_key_template.cc b/cc/experimental/pqcrypto/signature/dilithium_key_template.cc
index 933d06cce..fc431128e 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_key_template.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_key_template.cc
@@ -22,26 +22,51 @@
 #include "proto/tink.pb.h"
 #include "proto/tink.proto.h"
 
+extern "C" {
+#include "third_party/pqclean/crypto_sign/dilithium2/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium3/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium5/avx2/api.h"
+}
+
 namespace crypto {
 namespace tink {
 namespace {
 
+using google::crypto::tink::DilithiumKeyFormat;
 using google::crypto::tink::DilithiumPrivateKey;
 using google::crypto::tink::KeyTemplate;
 using google::crypto::tink::OutputPrefixType;
 
-KeyTemplate* NewDilithiumKeyTemplate() {
+KeyTemplate* NewDilithiumKeyTemplate(int32 key_size) {
   KeyTemplate* key_template = new KeyTemplate;
   key_template->set_type_url(
       absl::StrCat(kTypeGoogleapisCom, DilithiumPrivateKey().GetTypeName()));
   key_template->set_output_prefix_type(OutputPrefixType::TINK);
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(key_size);
+  key_format.SerializeToString(key_template->mutable_value());
+
   return key_template;
 }
 
 }  // anonymous namespace
 
-const google::crypto::tink::KeyTemplate& DilithiumKeyTemplate() {
-  static const KeyTemplate* key_template = NewDilithiumKeyTemplate();
+const google::crypto::tink::KeyTemplate& Dilithium2KeyTemplate() {
+  static const KeyTemplate* key_template =
+      NewDilithiumKeyTemplate(PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES);
+  return *key_template;
+}
+
+const google::crypto::tink::KeyTemplate& Dilithium3KeyTemplate() {
+  static const KeyTemplate* key_template =
+      NewDilithiumKeyTemplate(PQCLEAN_DILITHIUM3_AVX2_CRYPTO_SECRETKEYBYTES);
+  return *key_template;
+}
+
+const google::crypto::tink::KeyTemplate& Dilithium5KeyTemplate() {
+  static const KeyTemplate* key_template =
+      NewDilithiumKeyTemplate(PQCLEAN_DILITHIUM5_AVX2_CRYPTO_SECRETKEYBYTES);
   return *key_template;
 }
 
diff --git a/cc/experimental/pqcrypto/signature/dilithium_key_template.h b/cc/experimental/pqcrypto/signature/dilithium_key_template.h
index dea39896e..316610075 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_key_template.h
+++ b/cc/experimental/pqcrypto/signature/dilithium_key_template.h
@@ -23,7 +23,11 @@ namespace crypto {
 namespace tink {
 
 // Returns a KeyTemplate that generates new instances of DilithiumPrivateKey.
-const google::crypto::tink::KeyTemplate& DilithiumKeyTemplate();
+const google::crypto::tink::KeyTemplate& Dilithium2KeyTemplate();
+
+const google::crypto::tink::KeyTemplate& Dilithium3KeyTemplate();
+
+const google::crypto::tink::KeyTemplate& Dilithium5KeyTemplate();
 
 }  // namespace tink
 }  // namespace crypto
diff --git a/cc/experimental/pqcrypto/signature/dilithium_key_template_test.cc b/cc/experimental/pqcrypto/signature/dilithium_key_template_test.cc
index 8968f93d8..9cd98ae7d 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_key_template_test.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_key_template_test.cc
@@ -24,6 +24,13 @@
 #include "tink/experimental/pqcrypto/signature/dilithium_verify_key_manager.h"
 #include "tink/util/test_matchers.h"
 #include "proto/tink.pb.h"
+#include "proto/tink.proto.h"
+
+extern "C" {
+#include "third_party/pqclean/crypto_sign/dilithium2/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium3/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium5/avx2/api.h"
+}
 
 namespace crypto {
 namespace tink {
@@ -35,46 +42,67 @@ using ::google::crypto::tink::DilithiumPrivateKey;
 using ::google::crypto::tink::KeyTemplate;
 using ::google::crypto::tink::OutputPrefixType;
 
-TEST(DilithiumKeyTemplateTest, CheckInitialization) {
+struct DilithiumKeyTemplateTestCase {
+  std::string test_name;
+  int32 key_size;
+  KeyTemplate key_template;
+};
+
+using DilithiumKeyTemplateTest =
+    testing::TestWithParam<DilithiumKeyTemplateTestCase>;
+
+TEST_P(DilithiumKeyTemplateTest, CheckDilithiumInitialization) {
   std::string type_url =
       "type.googleapis.com/google.crypto.tink.DilithiumPrivateKey";
-  const KeyTemplate& key_template = DilithiumKeyTemplate();
+  const KeyTemplate& key_template = GetParam().key_template;
 
   EXPECT_EQ(type_url, key_template.type_url());
   EXPECT_EQ(OutputPrefixType::TINK, key_template.output_prefix_type());
 }
 
-TEST(DilithiumKeyTemplateTest, ValidateKeyFormat) {
-  const KeyTemplate& key_template = DilithiumKeyTemplate();
+TEST_P(DilithiumKeyTemplateTest, ValidateKeyFormat) {
+  const DilithiumKeyTemplateTestCase& test_case = GetParam();
   DilithiumKeyFormat key_format;
-  EXPECT_TRUE(key_format.ParseFromString(key_template.value()));
+
+  key_format.set_key_size(test_case.key_size);
   EXPECT_THAT(DilithiumSignKeyManager().ValidateKeyFormat(key_format), IsOk());
+  EXPECT_TRUE(key_format.ParseFromString(test_case.key_template.value()));
 }
 
-TEST(DilithiumKeyTemplateTest, SameReference) {
-  const KeyTemplate& key_template = DilithiumKeyTemplate();
-  const KeyTemplate& key_template_2 = DilithiumKeyTemplate();
+TEST_P(DilithiumKeyTemplateTest, SameReference) {
+  const KeyTemplate& key_template = GetParam().key_template;
+  const KeyTemplate& key_template_2 = GetParam().key_template;
 
   EXPECT_EQ(&key_template, &key_template_2);
 }
 
-TEST(DilithiumKeyTemplateTest, KeyManagerCompatibility) {
-  const KeyTemplate& key_template = DilithiumKeyTemplate();
-
+TEST_P(DilithiumKeyTemplateTest, KeyManagerCompatibility) {
   DilithiumSignKeyManager sign_key_manager;
   DilithiumVerifyKeyManager verify_key_manager;
   std::unique_ptr<KeyManager<PublicKeySign>> key_manager =
       internal::MakePrivateKeyManager<PublicKeySign>(&sign_key_manager,
                                                      &verify_key_manager);
-  EXPECT_EQ(key_manager->get_key_type(), key_template.type_url());
-
   DilithiumKeyFormat key_format;
-  EXPECT_TRUE(key_format.ParseFromString(key_template.value()));
-  util::StatusOr<std::unique_ptr<portable_proto::MessageLite>> new_key_result =
+  const DilithiumKeyTemplateTestCase& test_case = GetParam();
+
+  key_format.set_key_size(test_case.key_size);
+  util::StatusOr<std::unique_ptr<portable_proto::MessageLite>> new_key_result2 =
       key_manager->get_key_factory().NewKey(key_format);
-  EXPECT_THAT(new_key_result.status(), IsOk());
+  EXPECT_THAT(new_key_result2.status(), IsOk());
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    DilithiumKeyTemplateTests, DilithiumKeyTemplateTest,
+    testing::ValuesIn<DilithiumKeyTemplateTestCase>(
+        {{"Dilithium2", PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES,
+          Dilithium2KeyTemplate()},
+         {"Dilithium3", PQCLEAN_DILITHIUM3_AVX2_CRYPTO_SECRETKEYBYTES,
+          Dilithium2KeyTemplate()},
+         {"Dilithium5", PQCLEAN_DILITHIUM5_AVX2_CRYPTO_SECRETKEYBYTES,
+          Dilithium2KeyTemplate()}}),
+    [](const testing::TestParamInfo<DilithiumKeyTemplateTest::ParamType>&
+           info) { return info.param.test_name; });
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto
diff --git a/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager.cc b/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager.cc
index 1b8301fc6..46087c07c 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager.cc
@@ -32,6 +32,8 @@
 
 extern "C" {
 #include "third_party/pqclean/crypto_sign/dilithium2/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium3/avx2/api.h"
+#include "third_party/pqclean/crypto_sign/dilithium5/avx2/api.h"
 }
 
 namespace crypto {
@@ -43,21 +45,26 @@ using ::crypto::tink::util::Status;
 using ::crypto::tink::util::StatusOr;
 using ::google::crypto::tink::DilithiumKeyFormat;
 using ::google::crypto::tink::DilithiumPrivateKey;
+using ::google::crypto::tink::DilithiumPublicKey;
 
 StatusOr<DilithiumPrivateKey> DilithiumSignKeyManager::CreateKey(
     const DilithiumKeyFormat& key_format) const {
+  DilithiumPrivateKey dilithium_sk;
+  dilithium_sk.set_version(get_version());
+
+  DilithiumPublicKey* dilithium_pk = dilithium_sk.mutable_public_key();
+  dilithium_pk->set_version(get_version());
+
   util::StatusOr<
       std::pair<DilithiumPrivateKeyPqclean, DilithiumPublicKeyPqclean>>
-      key_pair = DilithiumPrivateKeyPqclean::GenerateKeyPair(
-          PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES);
+      key_pair =
+          DilithiumPrivateKeyPqclean::GenerateKeyPair(key_format.key_size());
+  if (!key_pair.status().ok()) {
+    return key_pair.status();
+  }
 
-  DilithiumPrivateKey dilithium_sk;
-  dilithium_sk.set_version(get_version());
   dilithium_sk.set_key_value(
       util::SecretDataAsStringView(key_pair->first.GetKeyData()));
-
-  auto dilithium_pk = dilithium_sk.mutable_public_key();
-  dilithium_pk->set_version(get_version());
   dilithium_pk->set_key_value(key_pair->second.GetKeyData());
 
   return dilithium_sk;
@@ -82,9 +89,13 @@ Status DilithiumSignKeyManager::ValidateKey(
   Status status = ValidateVersion(key.version(), get_version());
   if (!status.ok()) return status;
   if (key.key_value().length() !=
-      PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES) {
+          PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES &&
+      key.key_value().length() !=
+          PQCLEAN_DILITHIUM3_AVX2_CRYPTO_SECRETKEYBYTES &&
+      key.key_value().length() !=
+          PQCLEAN_DILITHIUM5_AVX2_CRYPTO_SECRETKEYBYTES) {
     return Status(util::error::INVALID_ARGUMENT,
-                  "The dilithium private key must be 2528-bytes long.");
+                  "Invalid dilithium private key size.");
   }
 
   return Status::OK;
diff --git a/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager_test.cc b/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager_test.cc
index cc3c02da0..ad15e8b33 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager_test.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_sign_key_manager_test.cc
@@ -37,6 +37,7 @@ extern "C" {
 
 namespace crypto {
 namespace tink {
+namespace {
 
 using ::crypto::tink::subtle::DilithiumPublicKeyPqclean;
 using ::crypto::tink::test::IsOk;
@@ -49,7 +50,13 @@ using ::testing::Eq;
 using ::testing::Not;
 using ::testing::SizeIs;
 
-namespace {
+struct DilithiumTestCase {
+  std::string test_name;
+  int private_key_size = 0;
+  int public_key_size = 0;
+};
+
+using DilithiumSignKeyManagerTest = testing::TestWithParam<DilithiumTestCase>;
 
 TEST(DilithiumSignKeyManagerTest, Basic) {
   EXPECT_THAT(DilithiumSignKeyManager().get_version(), Eq(0));
@@ -64,49 +71,74 @@ TEST(DilithiumSignKeyManagerTest, ValidateKeyFormat) {
               IsOk());
 }
 
-TEST(DilithiumSignKeyManagerTest, PrivateKeyWrongVersion) {
+TEST_P(DilithiumSignKeyManagerTest, PrivateKeyWrongVersion) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
+
   private_key->set_version(1);
   EXPECT_THAT(DilithiumSignKeyManager().ValidateKey(*private_key), Not(IsOk()));
 }
 
-TEST(DilithiumSignKeyManagerTest, CreateKey) {
+TEST_P(DilithiumSignKeyManagerTest, CreateKey) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
 
   EXPECT_THAT(private_key->version(), Eq(0));
   EXPECT_THAT(private_key->public_key().version(), Eq(private_key->version()));
-  EXPECT_THAT(private_key->key_value(),
-              SizeIs(PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES));
+  EXPECT_THAT(private_key->key_value(), SizeIs(test_case.private_key_size));
   EXPECT_THAT(private_key->public_key().key_value(),
-              SizeIs(PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES));
+              SizeIs(test_case.public_key_size));
 }
 
-TEST(DilithiumSignKeyManagerTest, CreateKeyValid) {
+TEST_P(DilithiumSignKeyManagerTest, CreateKeyValid) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
   EXPECT_THAT(DilithiumSignKeyManager().ValidateKey(*private_key), IsOk());
 }
 
-TEST(DilithiumSignKeyManagerTest, CreateKeyAlwaysNew) {
+TEST_P(DilithiumSignKeyManagerTest, CreateKeyAlwaysNew) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   absl::flat_hash_set<std::string> keys;
   int num_tests = 100;
   for (int i = 0; i < num_tests; ++i) {
     StatusOr<DilithiumPrivateKey> private_key =
-        DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+        DilithiumSignKeyManager().CreateKey(key_format);
     ASSERT_THAT(private_key.status(), IsOk());
     keys.insert(private_key->key_value());
   }
   EXPECT_THAT(keys, SizeIs(num_tests));
 }
 
-TEST(DilithiumSignKeyManagerTest, GetPublicKey) {
+TEST_P(DilithiumSignKeyManagerTest, GetPublicKey) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
 
   StatusOr<DilithiumPublicKey> public_key_or =
@@ -119,9 +151,14 @@ TEST(DilithiumSignKeyManagerTest, GetPublicKey) {
               Eq(private_key->public_key().key_value()));
 }
 
-TEST(DilithiumSignKeyManagerTest, Create) {
+TEST_P(DilithiumSignKeyManagerTest, Create) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   util::StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
 
   util::StatusOr<std::unique_ptr<PublicKeySign>> signer =
@@ -142,17 +179,21 @@ TEST(DilithiumSignKeyManagerTest, Create) {
   EXPECT_THAT((*verifier)->Verify(*signature, message), IsOk());
 }
 
-TEST(DilithiumSignKeyManagerTest, CreateDifferentKey) {
+TEST_P(DilithiumSignKeyManagerTest, CreateDifferentKey) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(test_case.private_key_size);
+
   util::StatusOr<DilithiumPrivateKey> private_key =
-      DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+      DilithiumSignKeyManager().CreateKey(key_format);
   ASSERT_THAT(private_key.status(), IsOk());
 
   util::StatusOr<std::unique_ptr<PublicKeySign>> signer =
       DilithiumSignKeyManager().GetPrimitive<PublicKeySign>(*private_key);
   ASSERT_THAT(signer.status(), IsOk());
 
-  std::string bad_public_key_data(PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES,
-                                  '@');
+  std::string bad_public_key_data(test_case.public_key_size, '@');
   util::StatusOr<DilithiumPublicKeyPqclean> dilithium_public_key =
       DilithiumPublicKeyPqclean::NewPublicKey(bad_public_key_data);
   util::StatusOr<std::unique_ptr<PublicKeyVerify>> verifier =
@@ -165,6 +206,19 @@ TEST(DilithiumSignKeyManagerTest, CreateDifferentKey) {
   EXPECT_THAT((*verifier)->Verify(*signature, message), Not(IsOk()));
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    DilithiumSignKeyManagerTests, DilithiumSignKeyManagerTest,
+    testing::ValuesIn<DilithiumTestCase>({
+        {"Dilithium2", PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES},
+        {"Dilithium3", PQCLEAN_DILITHIUM3_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM3_AVX2_CRYPTO_PUBLICKEYBYTES},
+        {"Dilithium5", PQCLEAN_DILITHIUM5_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM5_AVX2_CRYPTO_PUBLICKEYBYTES},
+    }),
+    [](const testing::TestParamInfo<DilithiumSignKeyManagerTest::ParamType>&
+           info) { return info.param.test_name; });
+
 }  // namespace
 
 }  // namespace tink
diff --git a/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager.cc b/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager.cc
index d313ca54b..134403e1a 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager.cc
@@ -61,9 +61,13 @@ Status DilithiumVerifyKeyManager::ValidateKey(
   if (!status.ok()) return status;
 
   if (key.key_value().length() !=
-      PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES) {
+          PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES &&
+      key.key_value().length() !=
+          PQCLEAN_DILITHIUM3_AVX2_CRYPTO_PUBLICKEYBYTES &&
+      key.key_value().length() !=
+          PQCLEAN_DILITHIUM5_AVX2_CRYPTO_PUBLICKEYBYTES) {
     return Status(util::error::INVALID_ARGUMENT,
-                  "The dilithium avx2 public key must be 1312-bytes long.");
+                  "Invalid dilithium public key size.");
   }
   return Status::OK;
 }
diff --git a/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager_test.cc b/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager_test.cc
index 00320311f..478577063 100644
--- a/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager_test.cc
+++ b/cc/experimental/pqcrypto/signature/dilithium_verify_key_manager_test.cc
@@ -38,6 +38,7 @@ extern "C" {
 
 namespace crypto {
 namespace tink {
+namespace {
 
 using ::crypto::tink::subtle::DilithiumPrivateKeyPqclean;
 using ::crypto::tink::test::IsOk;
@@ -49,16 +50,26 @@ using ::google::crypto::tink::KeyData;
 using ::testing::Eq;
 using ::testing::Not;
 
-namespace {
+struct DilithiumTestCase {
+  std::string test_name;
+  int private_key_size = 0;
+  int public_key_size = 0;
+};
+
+using DilithiumVerifyKeyManagerTest = testing::TestWithParam<DilithiumTestCase>;
 
 // Helper function that returns a valid dilithium private key.
-StatusOr<DilithiumPrivateKey> CreateValidPrivateKey() {
-  return DilithiumSignKeyManager().CreateKey(DilithiumKeyFormat());
+StatusOr<DilithiumPrivateKey> CreateValidPrivateKey(int32 private_key_size) {
+  DilithiumKeyFormat key_format;
+  key_format.set_key_size(private_key_size);
+
+  return DilithiumSignKeyManager().CreateKey(key_format);
 }
 
 // Helper function that returns a valid dilithium public key.
-StatusOr<DilithiumPublicKey> CreateValidPublicKey() {
-  StatusOr<DilithiumPrivateKey> private_key = CreateValidPrivateKey();
+StatusOr<DilithiumPublicKey> CreateValidPublicKey(int32 private_key_size) {
+  StatusOr<DilithiumPrivateKey> private_key =
+      CreateValidPrivateKey(private_key_size);
 
   if (!private_key.ok()) return private_key.status();
   return DilithiumSignKeyManager().GetPublicKey(*private_key);
@@ -77,15 +88,21 @@ TEST(DilithiumVerifyKeyManagerTest, ValidateEmptyKey) {
               Not(IsOk()));
 }
 
-TEST(DilithiumVerifyKeyManagerTest, PublicKeyValid) {
-  StatusOr<DilithiumPublicKey> public_key = CreateValidPublicKey();
+TEST_P(DilithiumVerifyKeyManagerTest, PublicKeyValid) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  StatusOr<DilithiumPublicKey> public_key =
+      CreateValidPublicKey(test_case.private_key_size);
   ASSERT_THAT(public_key.status(), IsOk());
 
   EXPECT_THAT(DilithiumVerifyKeyManager().ValidateKey(*public_key), IsOk());
 }
 
-TEST(DilithiumVerifyKeyManagerTest, PublicKeyWrongVersion) {
-  StatusOr<DilithiumPublicKey> public_key = CreateValidPublicKey();
+TEST_P(DilithiumVerifyKeyManagerTest, PublicKeyWrongVersion) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  StatusOr<DilithiumPublicKey> public_key =
+      CreateValidPublicKey(test_case.private_key_size);
   ASSERT_THAT(public_key.status(), IsOk());
 
   public_key->set_version(1);
@@ -93,8 +110,11 @@ TEST(DilithiumVerifyKeyManagerTest, PublicKeyWrongVersion) {
               Not(IsOk()));
 }
 
-TEST(DilithiumVerifyKeyManagerTest, PublicKeyWrongKeyLength) {
-  StatusOr<DilithiumPublicKey> public_key = CreateValidPublicKey();
+TEST_P(DilithiumVerifyKeyManagerTest, PublicKeyWrongKeyLength) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  StatusOr<DilithiumPublicKey> public_key =
+      CreateValidPublicKey(test_case.private_key_size);
   ASSERT_THAT(public_key.status(), IsOk());
 
   for (int keysize = 0; keysize < PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES;
@@ -105,8 +125,11 @@ TEST(DilithiumVerifyKeyManagerTest, PublicKeyWrongKeyLength) {
   }
 }
 
-TEST(DilithiumVerifyKeyManagerTest, Create) {
-  StatusOr<DilithiumPrivateKey> private_key = CreateValidPrivateKey();
+TEST_P(DilithiumVerifyKeyManagerTest, Create) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  StatusOr<DilithiumPrivateKey> private_key =
+      CreateValidPrivateKey(test_case.private_key_size);
   ASSERT_THAT(private_key.status(), IsOk());
 
   StatusOr<DilithiumPublicKey> public_key =
@@ -132,12 +155,16 @@ TEST(DilithiumVerifyKeyManagerTest, Create) {
   EXPECT_THAT((*verifier)->Verify(*signature, message), IsOk());
 }
 
-TEST(DilithiumVerifyKeyManagerTest, CreateDifferentPublicKey) {
-  StatusOr<DilithiumPrivateKey> private_key = CreateValidPrivateKey();
+TEST_P(DilithiumVerifyKeyManagerTest, CreateDifferentPublicKey) {
+  const DilithiumTestCase& test_case = GetParam();
+
+  StatusOr<DilithiumPrivateKey> private_key =
+      CreateValidPrivateKey(test_case.private_key_size);
   ASSERT_THAT(private_key.status(), IsOk());
 
   // Create a new public key derived from a diffferent private key.
-  StatusOr<DilithiumPrivateKey> new_private_key = CreateValidPrivateKey();
+  StatusOr<DilithiumPrivateKey> new_private_key =
+      CreateValidPrivateKey(test_case.private_key_size);
   ASSERT_THAT(new_private_key.status(), IsOk());
   StatusOr<DilithiumPublicKey> public_key =
       DilithiumSignKeyManager().GetPublicKey(*new_private_key);
@@ -162,6 +189,19 @@ TEST(DilithiumVerifyKeyManagerTest, CreateDifferentPublicKey) {
   EXPECT_THAT((*verifier)->Verify(*signature, message), Not(IsOk()));
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    DilithiumVerifyKeyManagerTests, DilithiumVerifyKeyManagerTest,
+    testing::ValuesIn<DilithiumTestCase>({
+        {"Dilithium2", PQCLEAN_DILITHIUM2_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM2_AVX2_CRYPTO_PUBLICKEYBYTES},
+        {"Dilithium3", PQCLEAN_DILITHIUM3_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM3_AVX2_CRYPTO_PUBLICKEYBYTES},
+        {"Dilithium5", PQCLEAN_DILITHIUM5_AVX2_CRYPTO_SECRETKEYBYTES,
+         PQCLEAN_DILITHIUM5_AVX2_CRYPTO_PUBLICKEYBYTES},
+    }),
+    [](const testing::TestParamInfo<DilithiumVerifyKeyManagerTest::ParamType>&
+           info) { return info.param.test_name; });
+
 }  // namespace
 
 }  // namespace tink