Add KMS envelope AEAD tests with other dek templates.

PiperOrigin-RevId: 541906234

4 files changed, 106 insertions, 0 deletions

diff --git a/cc/aead/BUILD.bazel b/cc/aead/BUILD.bazel
index 3d2cc49df..48ae6c5e1 100644
--- a/cc/aead/BUILD.bazel
+++ b/cc/aead/BUILD.bazel
@@ -681,6 +681,7 @@ cc_test(
         "//:aead",
         "//:keyset_handle",
         "//:registry",
+        "//internal:ssl_util",
         "//mac:mac_key_templates",
         "//proto:aes_gcm_cc_proto",
         "//util:status",
@@ -700,6 +701,7 @@ cc_test(
     size = "small",
     srcs = ["kms_envelope_aead_key_manager_test.cc"],
     deps = [
+        ":aead_config",
         ":aead_key_templates",
         ":aes_eax_key_manager",
         ":kms_envelope_aead",
@@ -711,6 +713,7 @@ cc_test(
         "//proto:kms_envelope_cc_proto",
         "//proto:tink_cc_proto",
         "//subtle:aead_test_util",
+        "//util:fake_kms_client",
         "//util:status",
         "//util:statusor",
         "//util:test_matchers",
diff --git a/cc/aead/CMakeLists.txt b/cc/aead/CMakeLists.txt
index 8b8b4f36f..fe45fef97 100644
--- a/cc/aead/CMakeLists.txt
+++ b/cc/aead/CMakeLists.txt
@@ -648,6 +648,7 @@ tink_cc_test(
     tink::core::aead
     tink::core::keyset_handle
     tink::core::registry
+    tink::internal::ssl_util
     tink::mac::mac_key_templates
     tink::util::status
     tink::util::statusor
@@ -661,6 +662,7 @@ tink_cc_test(
   SRCS
     kms_envelope_aead_key_manager_test.cc
   DEPS
+    tink::aead::aead_config
     tink::aead::aead_key_templates
     tink::aead::aes_eax_key_manager
     tink::aead::kms_envelope_aead
@@ -673,6 +675,7 @@ tink_cc_test(
     tink::core::kms_clients
     tink::core::registry
     tink::subtle::aead_test_util
+    tink::util::fake_kms_client
     tink::util::status
     tink::util::statusor
     tink::util::test_matchers
diff --git a/cc/aead/kms_envelope_aead_key_manager_test.cc b/cc/aead/kms_envelope_aead_key_manager_test.cc
index 865804a94..ad24fdb7e 100644
--- a/cc/aead/kms_envelope_aead_key_manager_test.cc
+++ b/cc/aead/kms_envelope_aead_key_manager_test.cc
@@ -26,6 +26,7 @@
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "tink/aead.h"
+#include "tink/aead/aead_config.h"
 #include "tink/aead/aead_key_templates.h"
 #include "tink/aead/aes_eax_key_manager.h"
 #include "tink/aead/kms_envelope_aead.h"
@@ -33,6 +34,7 @@
 #include "tink/kms_clients.h"
 #include "tink/registry.h"
 #include "tink/subtle/aead_test_util.h"
+#include "tink/util/fake_kms_client.h"
 #include "tink/util/status.h"
 #include "tink/util/statusor.h"
 #include "tink/util/test_matchers.h"
@@ -46,7 +48,9 @@ namespace tink {
 using ::crypto::tink::test::DummyAead;
 using ::crypto::tink::test::DummyKmsClient;
 using ::crypto::tink::test::IsOk;
+using ::crypto::tink::test::IsOkAndHolds;
 using ::crypto::tink::test::StatusIs;
+using ::google::crypto::tink::KeyTemplate;
 using ::google::crypto::tink::KmsEnvelopeAeadKey;
 using ::google::crypto::tink::KmsEnvelopeAeadKeyFormat;
 using ::testing::Eq;
@@ -228,6 +232,54 @@ TEST_F(KmsEnvelopeAeadKeyManagerCreateTest, CreateAeadUnboundKey) {
               IsOk());
 }
 
+class KmsEnvelopeAeadKeyManagerDekTemplatesTest
+    : public testing::TestWithParam<KeyTemplate> {
+  void SetUp() override { ASSERT_THAT(AeadConfig::Register(), IsOk()); }
+};
+
+TEST_P(KmsEnvelopeAeadKeyManagerDekTemplatesTest, EncryptDecryp) {
+  util::StatusOr<std::string> kek_uri_result =
+      test::FakeKmsClient::CreateFakeKeyUri();
+  ASSERT_THAT(kek_uri_result, IsOk());
+  std::string kek_uri = kek_uri_result.value();
+  util::Status register_fake_kms_client_status =
+      test::FakeKmsClient::RegisterNewClient(kek_uri, /*credentials_path=*/"");
+  ASSERT_THAT(register_fake_kms_client_status, IsOk());
+
+  KeyTemplate dek_template = GetParam();
+  KeyTemplate env_template =
+      AeadKeyTemplates::KmsEnvelopeAead(kek_uri, dek_template);
+  util::StatusOr<std::unique_ptr<KeysetHandle>> handle =
+      KeysetHandle::GenerateNew(env_template);
+  ASSERT_THAT(handle, IsOk());
+  util::StatusOr<std::unique_ptr<Aead>> envelope_aead =
+      (*handle)->GetPrimitive<Aead>();
+  ASSERT_THAT(envelope_aead, IsOk());
+
+  std::string plaintext = "plaintext";
+  std::string associated_data = "associated_data";
+  util::StatusOr<std::string> ciphertext =
+      (*envelope_aead)->Encrypt(plaintext, associated_data);
+  ASSERT_THAT(ciphertext, IsOk());
+  util::StatusOr<std::string> decrypted =
+      (*envelope_aead)->Decrypt(ciphertext.value(), associated_data);
+  EXPECT_THAT(decrypted, IsOkAndHolds(plaintext));
+
+  std::string invalid_associated_data = "invalid_associated_data";
+  util::StatusOr<std::string> decrypted_with_invalid_associated_data =
+      (*envelope_aead)->Decrypt(ciphertext.value(), invalid_associated_data);
+  EXPECT_THAT(decrypted_with_invalid_associated_data.status(), Not(IsOk()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    KmsEnvelopeAeadKeyManagerDekTemplatesTest,
+    KmsEnvelopeAeadKeyManagerDekTemplatesTest,
+    testing::Values(AeadKeyTemplates::Aes128Gcm(),
+                    AeadKeyTemplates::Aes256Gcm(),
+                    AeadKeyTemplates::Aes128CtrHmacSha256(),
+                    AeadKeyTemplates::Aes128Eax(),
+                    AeadKeyTemplates::Aes128GcmNoPrefix()));
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto
diff --git a/cc/aead/kms_envelope_aead_test.cc b/cc/aead/kms_envelope_aead_test.cc
index 83440db4d..49064618b 100644
--- a/cc/aead/kms_envelope_aead_test.cc
+++ b/cc/aead/kms_envelope_aead_test.cc
@@ -40,6 +40,7 @@
 #include "tink/util/test_matchers.h"
 #include "tink/util/test_util.h"
 #include "proto/aes_gcm.pb.h"
+#include "tink/internal/ssl_util.h"
 
 namespace crypto {
 namespace tink {
@@ -235,6 +236,53 @@ TEST_F(KmsEnvelopeAeadTest, MultipleEncryptionsProduceDifferentDeks) {
   }
 }
 
+class KmsEnvelopeAeadDekTemplatesTest
+    : public testing::TestWithParam<KeyTemplate> {
+  void SetUp() override { ASSERT_THAT(AeadConfig::Register(), IsOk()); }
+};
+
+TEST_P(KmsEnvelopeAeadDekTemplatesTest, EncryptDecrypt) {
+  // Use an AES-128-GCM primitive as the remote AEAD.
+  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm());
+  ASSERT_THAT(keyset_handle, IsOk());
+  util::StatusOr<std::unique_ptr<Aead>> remote_aead =
+      (*keyset_handle)->GetPrimitive<Aead>();
+
+  KeyTemplate dek_template = GetParam();
+  util::StatusOr<std::unique_ptr<Aead>> envelope_aead =
+      KmsEnvelopeAead::New(dek_template, *std::move(remote_aead));
+  ASSERT_THAT(envelope_aead, IsOk());
+
+  std::string plaintext = "plaintext";
+  std::string associated_data = "associated_data";
+  util::StatusOr<std::string> ciphertext =
+      (*envelope_aead)->Encrypt(plaintext, associated_data);
+  ASSERT_THAT(ciphertext, IsOk());
+  util::StatusOr<std::string> decrypted =
+      (*envelope_aead)->Decrypt(ciphertext.value(), associated_data);
+  EXPECT_THAT(decrypted, IsOkAndHolds(plaintext));
+}
+
+std::vector<KeyTemplate> GetTestTemplates() {
+  std::vector<KeyTemplate> templates = {
+    AeadKeyTemplates::Aes128Gcm(),
+    AeadKeyTemplates::Aes256Gcm(),
+    AeadKeyTemplates::Aes128CtrHmacSha256(),
+    AeadKeyTemplates::Aes128Eax(),
+    AeadKeyTemplates::Aes128GcmNoPrefix()
+  };
+  if (internal::IsBoringSsl()) {
+    templates.push_back(AeadKeyTemplates::XChaCha20Poly1305());
+    templates.push_back(AeadKeyTemplates::Aes256GcmSiv());
+  }
+  return templates;
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    KmsEnvelopeAeadDekTemplatesTest, KmsEnvelopeAeadDekTemplatesTest,
+    testing::ValuesIn(GetTestTemplates()));
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto