Add monitoring to deterministic AEAD wrapper.

If a MonitoringClientFactory is registered, then we create the DeterministicAeadWrapper with the corresponding MonitoringClients to report events from the EncryptDeterministically and DecryptDeterministically APIs.

PiperOrigin-RevId: 454808935

4 files changed, 362 insertions, 8 deletions

diff --git a/cc/daead/BUILD.bazel b/cc/daead/BUILD.bazel
index ffe890b49..2a4c93330 100644
--- a/cc/daead/BUILD.bazel
+++ b/cc/daead/BUILD.bazel
@@ -36,7 +36,10 @@ cc_library(
         "//:deterministic_aead",
         "//:primitive_set",
         "//:primitive_wrapper",
+        "//internal:monitoring_util",
+        "//internal:registry_impl",
         "//internal:util",
+        "//monitoring",
         "//proto:tink_cc_proto",
         "//util:status",
         "//util:statusor",
@@ -120,11 +123,16 @@ cc_test(
     srcs = ["deterministic_aead_wrapper_test.cc"],
     deps = [
         ":deterministic_aead_wrapper",
+        ":failing_daead",
         "//:deterministic_aead",
         "//:primitive_set",
+        "//internal:registry_impl",
+        "//monitoring",
+        "//monitoring:monitoring_client_mocks",
         "//util:status",
         "//util:test_matchers",
         "//util:test_util",
+        "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/cc/daead/CMakeLists.txt b/cc/daead/CMakeLists.txt
index 6aca5a1b9..4876217a2 100644
--- a/cc/daead/CMakeLists.txt
+++ b/cc/daead/CMakeLists.txt
@@ -37,7 +37,10 @@ tink_cc_library(
     tink::core::deterministic_aead
     tink::core::primitive_set
     tink::core::primitive_wrapper
+    tink::internal::monitoring_util
+    tink::internal::registry_impl
     tink::internal::util
+    tink::monitoring::monitoring
     tink::util::status
     tink::util::statusor
     tink::proto::tink_cc_proto
@@ -112,9 +115,14 @@ tink_cc_test(
     deterministic_aead_wrapper_test.cc
   DEPS
     tink::daead::deterministic_aead_wrapper
+    tink::daead::failing_daead
     gmock
+    absl::status
     tink::core::deterministic_aead
     tink::core::primitive_set
+    tink::internal::registry_impl
+    tink::monitoring::monitoring
+    tink::monitoring::monitoring_client_mocks
     tink::util::status
     tink::util::test_matchers
     tink::util::test_util
diff --git a/cc/daead/deterministic_aead_wrapper.cc b/cc/daead/deterministic_aead_wrapper.cc
index a36ae8ee2..3db1e99c1 100644
--- a/cc/daead/deterministic_aead_wrapper.cc
+++ b/cc/daead/deterministic_aead_wrapper.cc
@@ -22,7 +22,10 @@
 #include "absl/status/status.h"
 #include "tink/crypto_format.h"
 #include "tink/deterministic_aead.h"
+#include "tink/internal/monitoring_util.h"
+#include "tink/internal/registry_impl.h"
 #include "tink/internal/util.h"
+#include "tink/monitoring/monitoring.h"
 #include "tink/primitive_set.h"
 #include "tink/util/status.h"
 #include "tink/util/statusor.h"
@@ -32,6 +35,10 @@ namespace tink {
 
 namespace {
 
+constexpr absl::string_view kPrimitive = "daead";
+constexpr absl::string_view kEncryptApi = "encrypt";
+constexpr absl::string_view kDecryptApi = "decrypt";
+
 util::Status Validate(PrimitiveSet<DeterministicAead>* daead_set) {
   if (daead_set == nullptr) {
     return util::Status(absl::StatusCode::kInternal,
@@ -47,8 +54,13 @@ util::Status Validate(PrimitiveSet<DeterministicAead>* daead_set) {
 class  DeterministicAeadSetWrapper : public DeterministicAead {
  public:
   explicit DeterministicAeadSetWrapper(
-      std::unique_ptr<PrimitiveSet<DeterministicAead>> daead_set)
-      : daead_set_(std::move(daead_set)) {}
+      std::unique_ptr<PrimitiveSet<DeterministicAead>> daead_set,
+      std::unique_ptr<MonitoringClient> monitoring_encryption_client = nullptr,
+      std::unique_ptr<MonitoringClient> monitoring_decryption_client = nullptr)
+      : daead_set_(std::move(daead_set)),
+        monitoring_encryption_client_(std::move(monitoring_encryption_client)),
+        monitoring_decryption_client_(std::move(monitoring_decryption_client))
+        {}
 
   crypto::tink::util::StatusOr<std::string> EncryptDeterministically(
       absl::string_view plaintext,
@@ -62,6 +74,8 @@ class  DeterministicAeadSetWrapper : public DeterministicAead {
 
  private:
   std::unique_ptr<PrimitiveSet<DeterministicAead>> daead_set_;
+  std::unique_ptr<MonitoringClient> monitoring_encryption_client_;
+  std::unique_ptr<MonitoringClient> monitoring_decryption_client_;
 };
 
 util::StatusOr<std::string>
@@ -75,7 +89,16 @@ DeterministicAeadSetWrapper::EncryptDeterministically(
   auto encrypt_result =
       daead_set_->get_primary()->get_primitive().EncryptDeterministically(
           plaintext, associated_data);
-  if (!encrypt_result.ok()) return encrypt_result.status();
+  if (!encrypt_result.ok()) {
+    if (monitoring_encryption_client_ != nullptr) {
+      monitoring_encryption_client_->LogFailure();
+    }
+    return encrypt_result.status();
+  }
+  if (monitoring_encryption_client_ != nullptr) {
+    monitoring_encryption_client_->Log(daead_set_->get_primary()->get_key_id(),
+                                       plaintext.size());
+  }
   const std::string& key_id = daead_set_->get_primary()->get_identifier();
   return key_id + encrypt_result.value();
 }
@@ -94,11 +117,15 @@ DeterministicAeadSetWrapper::DecryptDeterministically(
     if (primitives_result.ok()) {
       absl::string_view raw_ciphertext =
           ciphertext.substr(CryptoFormat::kNonRawPrefixSize);
-      for (auto& daead_entry : *(primitives_result.value())) {
+      for (const auto& daead_entry : *(primitives_result.value())) {
         DeterministicAead& daead = daead_entry->get_primitive();
         auto decrypt_result =
             daead.DecryptDeterministically(raw_ciphertext, associated_data);
         if (decrypt_result.ok()) {
+          if (monitoring_decryption_client_ != nullptr) {
+            monitoring_decryption_client_->Log(daead_entry->get_key_id(),
+                                               raw_ciphertext.size());
+          }
           return std::move(decrypt_result.value());
         } else {
           // LOG that a matching key didn't decrypt the ciphertext.
@@ -110,15 +137,22 @@ DeterministicAeadSetWrapper::DecryptDeterministically(
   // No matching key succeeded with decryption, try all RAW keys.
   auto raw_primitives_result = daead_set_->get_raw_primitives();
   if (raw_primitives_result.ok()) {
-    for (auto& daead_entry : *(raw_primitives_result.value())) {
+    for (const auto& daead_entry : *(raw_primitives_result.value())) {
       DeterministicAead& daead = daead_entry->get_primitive();
       auto decrypt_result =
           daead.DecryptDeterministically(ciphertext, associated_data);
       if (decrypt_result.ok()) {
+        if (monitoring_decryption_client_ != nullptr) {
+          monitoring_decryption_client_->Log(daead_entry->get_key_id(),
+                                             ciphertext.size());
+        }
         return std::move(decrypt_result.value());
       }
     }
   }
+  if (monitoring_decryption_client_ != nullptr) {
+    monitoring_decryption_client_->LogFailure();
+  }
   return util::Status(absl::StatusCode::kInvalidArgument, "decryption failed");
 }
 
@@ -129,9 +163,39 @@ DeterministicAeadWrapper::Wrap(
     std::unique_ptr<PrimitiveSet<DeterministicAead>> primitive_set) const {
   util::Status status = Validate(primitive_set.get());
   if (!status.ok()) return status;
-  std::unique_ptr<DeterministicAead> daead(
-      new DeterministicAeadSetWrapper(std::move(primitive_set)));
-  return std::move(daead);
+
+  MonitoringClientFactory* const monitoring_factory =
+      internal::RegistryImpl::GlobalInstance().GetMonitoringClientFactory();
+
+  // Monitoring is not enabled. Create a wrapper without monitoring clients.
+  if (monitoring_factory == nullptr) {
+    return {absl::make_unique<DeterministicAeadSetWrapper>(
+        std::move(primitive_set))};
+  }
+
+  util::StatusOr<MonitoringKeySetInfo> keyset_info =
+      internal::MonitoringKeySetInfoFromPrimitiveSet(*primitive_set);
+  if (!keyset_info.ok()) {
+    return keyset_info.status();
+  }
+
+  util::StatusOr<std::unique_ptr<MonitoringClient>>
+      monitoring_encryption_client = monitoring_factory->New(
+          MonitoringContext(kPrimitive, kEncryptApi, *keyset_info));
+  if (!monitoring_encryption_client.ok()) {
+    return monitoring_encryption_client.status();
+  }
+
+  util::StatusOr<std::unique_ptr<MonitoringClient>>
+      monitoring_decryption_client = monitoring_factory->New(
+          MonitoringContext(kPrimitive, kDecryptApi, *keyset_info));
+  if (!monitoring_decryption_client.ok()) {
+    return monitoring_decryption_client.status();
+  }
+
+  return {absl::make_unique<DeterministicAeadSetWrapper>(
+      std::move(primitive_set), *std::move(monitoring_encryption_client),
+      *std::move(monitoring_decryption_client))};
 }
 
 }  // namespace tink
diff --git a/cc/daead/deterministic_aead_wrapper_test.cc b/cc/daead/deterministic_aead_wrapper_test.cc
index a34daac20..c4d983af8 100644
--- a/cc/daead/deterministic_aead_wrapper_test.cc
+++ b/cc/daead/deterministic_aead_wrapper_test.cc
@@ -20,8 +20,14 @@
 #include <string>
 #include <utility>
 
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/status/status.h"
+#include "tink/daead/failing_daead.h"
 #include "tink/deterministic_aead.h"
+#include "tink/internal/registry_impl.h"
+#include "tink/monitoring/monitoring.h"
+#include "tink/monitoring/monitoring_client_mocks.h"
 #include "tink/primitive_set.h"
 #include "tink/util/status.h"
 #include "tink/util/test_matchers.h"
@@ -29,9 +35,19 @@
 
 using ::crypto::tink::test::DummyDeterministicAead;
 using ::crypto::tink::test::IsOk;
+using ::crypto::tink::test::IsOkAndHolds;
+using ::crypto::tink::test::StatusIs;
 using ::google::crypto::tink::KeysetInfo;
 using ::google::crypto::tink::KeyStatusType;
 using ::google::crypto::tink::OutputPrefixType;
+using ::testing::_;
+using ::testing::ByMove;
+using ::testing::IsNull;
+using ::testing::NiceMock;
+using ::testing::Not;
+using ::testing::NotNull;
+using ::testing::Return;
+using ::testing::Test;
 
 namespace crypto {
 namespace tink {
@@ -139,6 +155,264 @@ TEST_F(DeterministicAeadSetWrapperTest, testBasic) {
   }
 }
 
+KeysetInfo::KeyInfo PopulateKeyInfo(uint32_t key_id,
+                                    OutputPrefixType out_prefix_type,
+                                    KeyStatusType status) {
+  KeysetInfo::KeyInfo key_info;
+  key_info.set_output_prefix_type(out_prefix_type);
+  key_info.set_key_id(key_id);
+  key_info.set_status(status);
+  return key_info;
+}
+
+// Creates a test keyset info object.
+KeysetInfo CreateTestKeysetInfo() {
+  KeysetInfo keyset_info;
+  *keyset_info.add_key_info() =
+      PopulateKeyInfo(/*key_id=*/1234543, OutputPrefixType::TINK,
+                      /*status=*/KeyStatusType::ENABLED);
+  *keyset_info.add_key_info() =
+      PopulateKeyInfo(/*key_id=*/726329, OutputPrefixType::LEGACY,
+                      /*status=*/KeyStatusType::ENABLED);
+  *keyset_info.add_key_info() =
+      PopulateKeyInfo(/*key_id=*/7213743, OutputPrefixType::TINK,
+                      /*status=*/KeyStatusType::ENABLED);
+  return keyset_info;
+}
+
+// Tests for the monitoring behavior.
+class DeterministicAeadSetWrapperWithMonitoringTest : public Test {
+ protected:
+  // Perform some common initialization: reset the global registry, set expected
+  // calls for the mock monitoring factory and the returned clients.
+  void SetUp() override {
+    Registry::Reset();
+
+    // Setup mocks for catching Monitoring calls.
+    auto monitoring_client_factory =
+        absl::make_unique<MockMonitoringClientFactory>();
+    auto encryption_monitoring_client =
+        absl::make_unique<NiceMock<MockMonitoringClient>>();
+    encryption_monitoring_client_ = encryption_monitoring_client.get();
+    auto decryption_monitoring_client =
+        absl::make_unique<NiceMock<MockMonitoringClient>>();
+    decryption_monitoring_client_ = decryption_monitoring_client.get();
+
+    // Monitoring tests expect that the client factory will create the
+    // corresponding MockMonitoringClients.
+    EXPECT_CALL(*monitoring_client_factory, New(_))
+        .WillOnce(
+            Return(ByMove(util::StatusOr<std::unique_ptr<MonitoringClient>>(
+                std::move(encryption_monitoring_client)))))
+        .WillOnce(
+            Return(ByMove(util::StatusOr<std::unique_ptr<MonitoringClient>>(
+                std::move(decryption_monitoring_client)))));
+
+    ASSERT_THAT(internal::RegistryImpl::GlobalInstance()
+                    .RegisterMonitoringClientFactory(
+                        std::move(monitoring_client_factory)),
+                IsOk());
+    ASSERT_THAT(
+        internal::RegistryImpl::GlobalInstance().GetMonitoringClientFactory(),
+        Not(IsNull()));
+  }
+
+  // Cleanup the registry to avoid mock leaks.
+  ~DeterministicAeadSetWrapperWithMonitoringTest() override {
+    Registry::Reset();
+  }
+
+  MockMonitoringClient* encryption_monitoring_client_;
+  MockMonitoringClient* decryption_monitoring_client_;
+};
+
+// Test that successful encrypt operations are logged.
+TEST_F(DeterministicAeadSetWrapperWithMonitoringTest,
+       WrapKeysetWithMonitoringEncryptSuccess) {
+  // Create a primitive set and fill it with some entries
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+  const absl::flat_hash_map<std::string, std::string> annotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+  auto daead_primitive_set =
+      absl::make_unique<PrimitiveSet<DeterministicAead>>(annotations);
+  ASSERT_THAT(
+      daead_primitive_set
+          ->AddPrimitive(absl::make_unique<DummyDeterministicAead>("daead0"),
+                         keyset_info.key_info(0))
+          .status(), IsOk());
+  ASSERT_THAT(
+      daead_primitive_set
+          ->AddPrimitive(absl::make_unique<DummyDeterministicAead>("daead1"),
+                         keyset_info.key_info(1))
+          .status(), IsOk());
+  // Set the last as primary.
+  util::StatusOr<PrimitiveSet<DeterministicAead>::Entry<DeterministicAead>*>
+      last = daead_primitive_set->AddPrimitive(
+          absl::make_unique<DummyDeterministicAead>("daead2"),
+          keyset_info.key_info(2));
+  ASSERT_THAT(last.status(), IsOk());
+  ASSERT_THAT(daead_primitive_set->set_primary(*last), IsOk());
+  // Record the ID of the primary key.
+  const uint32_t primary_key_id = keyset_info.key_info(2).key_id();
+
+  // Create a deterministic AEAD and encrypt some data.
+  util::StatusOr<std::unique_ptr<DeterministicAead>> daead =
+      DeterministicAeadWrapper().Wrap(std::move(daead_primitive_set));
+  ASSERT_THAT(daead, IsOkAndHolds(NotNull()));
+
+  constexpr absl::string_view plaintext = "This is some plaintext!";
+  constexpr absl::string_view associated_data = "Some associated data!";
+
+  // Check that calling EncryptDeterministically triggers a Log() call.
+  EXPECT_CALL(*encryption_monitoring_client_,
+              Log(primary_key_id, plaintext.size()));
+  util::StatusOr<std::string> ciphertext =
+      (*daead)->EncryptDeterministically(plaintext, associated_data);
+  EXPECT_THAT(ciphertext.status(), IsOk());
+}
+
+// Test that successful encrypt operations are logged.
+TEST_F(DeterministicAeadSetWrapperWithMonitoringTest,
+       WrapKeysetWithMonitoringDecryptSuccess) {
+  // Create a primitive set and fill it with some entries
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+  const absl::flat_hash_map<std::string, std::string> annotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+  auto daead_primitive_set =
+      absl::make_unique<PrimitiveSet<DeterministicAead>>(annotations);
+  ASSERT_THAT(
+      daead_primitive_set
+          ->AddPrimitive(absl::make_unique<DummyDeterministicAead>("daead0"),
+                         keyset_info.key_info(0))
+          .status(),
+      IsOk());
+  ASSERT_THAT(
+      daead_primitive_set
+          ->AddPrimitive(absl::make_unique<DummyDeterministicAead>("daead1"),
+                         keyset_info.key_info(1))
+          .status(),
+      IsOk());
+  // Set the last as primary.
+  util::StatusOr<PrimitiveSet<DeterministicAead>::Entry<DeterministicAead>*>
+      last = daead_primitive_set->AddPrimitive(
+          absl::make_unique<DummyDeterministicAead>("daead2"),
+          keyset_info.key_info(2));
+  ASSERT_THAT(last.status(), IsOk());
+  ASSERT_THAT(daead_primitive_set->set_primary(*last), IsOk());
+  // Record the ID of the primary key.
+  const uint32_t primary_key_id = keyset_info.key_info(2).key_id();
+
+
+  // Create a deterministic AEAD and encrypt/decrypt some data.
+  util::StatusOr<std::unique_ptr<DeterministicAead>> daead =
+      DeterministicAeadWrapper().Wrap(std::move(daead_primitive_set));
+  ASSERT_THAT(daead, IsOkAndHolds(NotNull()));
+
+  constexpr absl::string_view plaintext = "This is some plaintext!";
+  constexpr absl::string_view associated_data = "Some associated data!";
+
+
+  // Check that calling DecryptDeterministically triggers a Log() call.
+  util::StatusOr<std::string> ciphertext =
+      (*daead)->EncryptDeterministically(plaintext, associated_data);
+  EXPECT_THAT(ciphertext.status(), IsOk());
+
+  // In the log expect the size of the ciphertext without the non-raw prefix.
+  EXPECT_CALL(*decryption_monitoring_client_,
+              Log(primary_key_id,
+                  ciphertext->size() - CryptoFormat::kNonRawPrefixSize));
+  EXPECT_THAT(
+      (*daead)->DecryptDeterministically(*ciphertext, associated_data).status(),
+      IsOk());
+}
+
+TEST_F(DeterministicAeadSetWrapperWithMonitoringTest,
+       WrapKeysetWithMonitoringEncryptFailures) {
+  // Create a primitive set and fill it with some entries.
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+  const absl::flat_hash_map<std::string, std::string> annotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+  auto daead_primitive_set =
+      absl::make_unique<PrimitiveSet<DeterministicAead>>(annotations);
+  ASSERT_THAT(daead_primitive_set
+                  ->AddPrimitive(CreateAlwaysFailingDeterministicAead("daead0"),
+                                 keyset_info.key_info(0))
+                  .status(),
+              IsOk());
+  ASSERT_THAT(daead_primitive_set
+                  ->AddPrimitive(CreateAlwaysFailingDeterministicAead("daead1"),
+                                 keyset_info.key_info(1))
+                  .status(),
+              IsOk());
+  // Set the last as primary.
+  util::StatusOr<PrimitiveSet<DeterministicAead>::Entry<DeterministicAead>*>
+      last = daead_primitive_set->AddPrimitive(
+          CreateAlwaysFailingDeterministicAead("daead2"),
+          keyset_info.key_info(2));
+  ASSERT_THAT(last.status(), IsOk());
+  ASSERT_THAT(daead_primitive_set->set_primary(*last), IsOk());
+
+
+  // Create a deterministic AEAD and encrypt.
+  util::StatusOr<std::unique_ptr<DeterministicAead>> daead =
+      DeterministicAeadWrapper().Wrap(std::move(daead_primitive_set));
+  ASSERT_THAT(daead, IsOkAndHolds(NotNull()));
+
+  constexpr absl::string_view plaintext = "This is some plaintext!";
+  constexpr absl::string_view associated_data = "Some associated data!";
+
+
+  // Check that calling EncryptDeterministically triggers a LogFailure() call.
+  EXPECT_CALL(*encryption_monitoring_client_, LogFailure());
+  util::StatusOr<std::string> ciphertext =
+      (*daead)->EncryptDeterministically(plaintext, associated_data);
+  EXPECT_THAT(ciphertext.status(), StatusIs(absl::StatusCode::kInternal));
+}
+
+// Test that monitoring logs decryption failures correctly.
+TEST_F(DeterministicAeadSetWrapperWithMonitoringTest,
+       WrapKeysetWithMonitoringDecryptFailures) {
+  // Create a primitive set and fill it with some entries.
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+  const absl::flat_hash_map<std::string, std::string> annotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+  auto daead_primitive_set =
+      absl::make_unique<PrimitiveSet<DeterministicAead>>(annotations);
+  ASSERT_THAT(daead_primitive_set
+                  ->AddPrimitive(CreateAlwaysFailingDeterministicAead("daead0"),
+                                 keyset_info.key_info(0))
+                  .status(),
+              IsOk());
+  ASSERT_THAT(daead_primitive_set
+                  ->AddPrimitive(CreateAlwaysFailingDeterministicAead("daead1"),
+                                 keyset_info.key_info(1))
+                  .status(),
+              IsOk());
+  // Set the last as primary.
+  util::StatusOr<PrimitiveSet<DeterministicAead>::Entry<DeterministicAead>*>
+      last = daead_primitive_set->AddPrimitive(
+          CreateAlwaysFailingDeterministicAead("daead2"),
+          keyset_info.key_info(2));
+  ASSERT_THAT(last.status(), IsOk());
+  ASSERT_THAT(daead_primitive_set->set_primary(*last), IsOk());
+
+
+  // Create a deterministic AEAD and decrypt.
+  util::StatusOr<std::unique_ptr<DeterministicAead>> daead =
+      DeterministicAeadWrapper().Wrap(std::move(daead_primitive_set));
+  ASSERT_THAT(daead, IsOkAndHolds(NotNull()));
+
+  constexpr absl::string_view associated_data = "Some associated data!";
+  constexpr absl::string_view ciphertext = "This is some ciphertext!";
+
+
+  // Check that calling DecryptDeterministically triggers a LogFailure() call.
+  EXPECT_CALL(*decryption_monitoring_client_, LogFailure());
+  EXPECT_THAT(
+      (*daead)->DecryptDeterministically(ciphertext, associated_data).status(),
+      StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto