Add support for monitoring AEAD encrypt/decrypt operations

 * Modify the `AeadWrapper::Wrap` method to create monitoring clients for the inner `AeadSetWrapper` class if monitoring is enabled (i.e., if a factory is registered with `RegistryImpl`).
 * Modify `AeadSetWrapper` to call `Log|LogFailure` when encryption/decryption succeed|fail.
 * Add unit tests to make sure logging calls are made when monitoring is enabled.

PiperOrigin-RevId: 439554406

4 files changed, 275 insertions, 32 deletions

diff --git a/cc/aead/BUILD.bazel b/cc/aead/BUILD.bazel
index 339cbde5d..ee46f4ca3 100644
--- a/cc/aead/BUILD.bazel
+++ b/cc/aead/BUILD.bazel
@@ -12,10 +12,14 @@ cc_library(
         "//:crypto_format",
         "//:primitive_set",
         "//:primitive_wrapper",
+        "//internal:monitoring_util",
+        "//internal:registry_impl",
         "//internal:util",
+        "//monitoring",
         "//proto:tink_cc_proto",
         "//util:status",
         "//util:statusor",
+        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
@@ -328,12 +332,22 @@ cc_test(
     srcs = ["aead_wrapper_test.cc"],
     deps = [
         ":aead_wrapper",
+        ":mock_aead",
         "//:aead",
+        "//:crypto_format",
         "//:primitive_set",
+        "//:registry",
+        "//internal:registry_impl",
+        "//monitoring",
+        "//monitoring:monitoring_client_mocks",
         "//proto:tink_cc_proto",
         "//util:status",
         "//util:test_matchers",
         "//util:test_util",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/cc/aead/CMakeLists.txt b/cc/aead/CMakeLists.txt
index 4ed14ec09..0daec4a8e 100644
--- a/cc/aead/CMakeLists.txt
+++ b/cc/aead/CMakeLists.txt
@@ -8,13 +8,17 @@ tink_cc_library(
     aead_wrapper.cc
     aead_wrapper.h
   DEPS
+    absl::memory
     absl::status
     absl::strings
     tink::core::aead
     tink::core::crypto_format
     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
@@ -305,9 +309,19 @@ tink_cc_test(
     aead_wrapper_test.cc
   DEPS
     tink::aead::aead_wrapper
+    tink::aead::mock_aead
     gmock
+    absl::flat_hash_map
+    absl::memory
+    absl::status
+    absl::strings
     tink::core::aead
+    tink::core::crypto_format
     tink::core::primitive_set
+    tink::core::registry
+    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/aead/aead_wrapper.cc b/cc/aead/aead_wrapper.cc
index 40223d451..d31823bcc 100644
--- a/cc/aead/aead_wrapper.cc
+++ b/cc/aead/aead_wrapper.cc
@@ -20,11 +20,15 @@
 #include <string>
 #include <utility>
 
+#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "tink/aead.h"
 #include "tink/crypto_format.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"
@@ -33,6 +37,10 @@ namespace crypto {
 namespace tink {
 namespace {
 
+constexpr absl::string_view kPrimitive = "aead";
+constexpr absl::string_view kEncryptApi = "encrypt";
+constexpr absl::string_view kDecryptApi = "decrypt";
+
 util::Status Validate(PrimitiveSet<Aead>* aead_set) {
   if (aead_set == nullptr) {
     return util::Status(absl::StatusCode::kInternal,
@@ -45,12 +53,17 @@ util::Status Validate(PrimitiveSet<Aead>* aead_set) {
   return util::OkStatus();
 }
 
+// The actual wrapper.
 class AeadSetWrapper : public Aead {
  public:
-  explicit AeadSetWrapper(std::unique_ptr<PrimitiveSet<Aead>> aead_set)
-      : aead_set_(std::move(aead_set)) {}
-
-  ~AeadSetWrapper() override {}
+  explicit AeadSetWrapper(
+      std::unique_ptr<PrimitiveSet<Aead>> aead_set,
+      std::unique_ptr<MonitoringClient> monitoring_encryption_client = nullptr,
+      std::unique_ptr<MonitoringClient> monitoring_decryption_client = nullptr)
+      : aead_set_(std::move(aead_set)),
+        monitoring_encryption_client_(std::move(monitoring_encryption_client)),
+        monitoring_decryption_client_(std::move(monitoring_decryption_client)) {
+  }
 
   util::StatusOr<std::string> Encrypt(
       absl::string_view plaintext,
@@ -62,20 +75,26 @@ class AeadSetWrapper : public Aead {
 
  private:
   std::unique_ptr<PrimitiveSet<Aead>> aead_set_;
+  std::unique_ptr<MonitoringClient> monitoring_encryption_client_;
+  std::unique_ptr<MonitoringClient> monitoring_decryption_client_;
 };
 
 util::StatusOr<std::string> AeadSetWrapper::Encrypt(
     absl::string_view plaintext, absl::string_view associated_data) const {
-  // BoringSSL expects a non-null pointer for plaintext and additional_data,
-  // regardless of whether the size is 0.
-  plaintext = internal::EnsureStringNonNull(plaintext);
   associated_data = internal::EnsureStringNonNull(associated_data);
   const Aead& primitive = aead_set_->get_primary()->get_primitive();
   util::StatusOr<std::string> ciphertext =
       primitive.Encrypt(plaintext, associated_data);
   if (!ciphertext.ok()) {
+    if (monitoring_encryption_client_ != nullptr) {
+      monitoring_encryption_client_->LogFailure();
+    }
     return ciphertext.status();
   }
+  if (monitoring_encryption_client_ != nullptr) {
+    monitoring_encryption_client_->Log(aead_set_->get_primary()->get_key_id(),
+                                       plaintext.size());
+  }
   const std::string& key_id = aead_set_->get_primary()->get_identifier();
   return absl::StrCat(key_id, *ciphertext);
 }
@@ -100,6 +119,10 @@ util::StatusOr<std::string> AeadSetWrapper::Decrypt(
         util::StatusOr<std::string> plaintext =
             aead.Decrypt(raw_ciphertext, associated_data);
         if (plaintext.ok()) {
+          if (monitoring_decryption_client_ != nullptr) {
+            monitoring_decryption_client_->Log(aead_entry->get_key_id(),
+                                               raw_ciphertext.size());
+          }
           return plaintext;
         }
       }
@@ -116,14 +139,21 @@ util::StatusOr<std::string> AeadSetWrapper::Decrypt(
       util::StatusOr<std::string> plaintext =
           aead.Decrypt(ciphertext, associated_data);
       if (plaintext.ok()) {
+        if (monitoring_decryption_client_ != nullptr) {
+          monitoring_decryption_client_->Log(aead_entry->get_key_id(),
+                                             ciphertext.size());
+        }
         return plaintext;
       }
     }
   }
+  if (monitoring_decryption_client_ != nullptr) {
+    monitoring_decryption_client_->LogFailure();
+  }
   return util::Status(absl::StatusCode::kInvalidArgument, "decryption failed");
 }
 
-}  // anonymous namespace
+}  // namespace
 
 util::StatusOr<std::unique_ptr<Aead>> AeadWrapper::Wrap(
     std::unique_ptr<PrimitiveSet<Aead>> aead_set) const {
@@ -131,8 +161,38 @@ util::StatusOr<std::unique_ptr<Aead>> AeadWrapper::Wrap(
   if (!status.ok()) {
     return status;
   }
-  std::unique_ptr<Aead> aead(new AeadSetWrapper(std::move(aead_set)));
-  return std::move(aead);
+
+  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<AeadSetWrapper>(std::move(aead_set))};
+  }
+
+  util::StatusOr<MonitoringKeySetInfo> keyset_info =
+      internal::MonitoringKeySetInfoFromPrimitiveSet(*aead_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<AeadSetWrapper>(
+      std::move(aead_set), *std::move(monitoring_encryption_client),
+      *std::move(monitoring_decryption_client))};
 }
 
 }  // namespace tink
diff --git a/cc/aead/aead_wrapper_test.cc b/cc/aead/aead_wrapper_test.cc
index 245f06819..b26e603ff 100644
--- a/cc/aead/aead_wrapper_test.cc
+++ b/cc/aead/aead_wrapper_test.cc
@@ -22,8 +22,18 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
 #include "tink/aead.h"
+#include "tink/aead/mock_aead.h"
+#include "tink/crypto_format.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/registry.h"
 #include "tink/util/status.h"
 #include "tink/util/test_matchers.h"
 #include "tink/util/test_util.h"
@@ -39,9 +49,15 @@ 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::HasSubstr;
+using ::testing::IsNull;
 using ::testing::IsSubstring;
 using ::testing::Not;
+using ::testing::Return;
+using ::testing::StrictMock;
+using ::testing::Test;
 
 void PopulateKeyInfo(KeysetInfo::KeyInfo* key_info, uint32_t key_id,
                      OutputPrefixType out_prefix_type, KeyStatusType status) {
@@ -50,6 +66,21 @@ void PopulateKeyInfo(KeysetInfo::KeyInfo* key_info, uint32_t key_id,
   key_info->set_status(status);
 }
 
+// Creates a test keyset info object.
+KeysetInfo CreateTestKeysetInfo() {
+  KeysetInfo keyset_info;
+  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/1234543,
+                  OutputPrefixType::TINK,
+                  /*status=*/KeyStatusType::ENABLED);
+  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/726329,
+                  OutputPrefixType::LEGACY,
+                  /*status=*/KeyStatusType::ENABLED);
+  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/7213743,
+                  OutputPrefixType::TINK,
+                  /*status=*/KeyStatusType::ENABLED);
+  return keyset_info;
+}
+
 TEST(AeadSetWrapperTest, WrapNullptr) {
   AeadWrapper wrapper;
   util::StatusOr<std::unique_ptr<Aead>> aead = wrapper.Wrap(nullptr);
@@ -70,17 +101,7 @@ TEST(AeadSetWrapperTest, WrapEmpty) {
 }
 
 TEST(AeadSetWrapperTest, Basic) {
-  KeysetInfo keyset_info;
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/1234543,
-                  OutputPrefixType::TINK,
-                  /*status=*/KeyStatusType::ENABLED);
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/726329,
-                  OutputPrefixType::LEGACY,
-                  /*status=*/KeyStatusType::ENABLED);
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/7213743,
-                  OutputPrefixType::TINK,
-                  /*status=*/KeyStatusType::ENABLED);
-
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
   std::string aead_name_0 = "aead0";
   std::string aead_name_1 = "aead1";
   std::string aead_name_2 = "aead2";
@@ -126,17 +147,7 @@ TEST(AeadSetWrapperTest, Basic) {
 }
 
 TEST(AeadSetWrapperTest, DecryptNonPrimary) {
-  KeysetInfo keyset_info;
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/1234543,
-                  OutputPrefixType::TINK,
-                  /*status=*/KeyStatusType::ENABLED);
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/726329,
-                  OutputPrefixType::LEGACY,
-                  /*status=*/KeyStatusType::ENABLED);
-  PopulateKeyInfo(keyset_info.add_key_info(), /*key_id=*/7213743,
-                  OutputPrefixType::TINK,
-                  /*status=*/KeyStatusType::ENABLED);
-
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
   std::string aead_name_0 = "aead0";
   std::string aead_name_1 = "aead1";
   std::string aead_name_2 = "aead2";
@@ -180,6 +191,150 @@ TEST(AeadSetWrapperTest, DecryptNonPrimary) {
       aead->Decrypt(complete_ciphertext, aad);
   EXPECT_THAT(decrypted_plaintext.status(), IsOk());
 }
+
+// Tests with monitoring enabled.
+class AeadSetWrapperTestWithMonitoring : 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();
+    auto monitoring_client_factory =
+        absl::make_unique<MockMonitoringClientFactory>();
+
+    auto encryption_monitoring_client =
+        absl::make_unique<StrictMock<MockMonitoringClient>>();
+    encryption_monitoring_client_ptr_ = encryption_monitoring_client.get();
+    auto decryption_monitoring_client =
+        absl::make_unique<StrictMock<MockMonitoringClient>>();
+    decryption_monitoring_client_ptr_ = decryption_monitoring_client.get();
+
+    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.
+  void TearDown() override { Registry::Reset(); }
+
+  MockMonitoringClient* encryption_monitoring_client_ptr_;
+  MockMonitoringClient* decryption_monitoring_client_ptr_;
+};
+
+// Test that successful encrypt/decrypt operations are logged.
+TEST_F(AeadSetWrapperTestWithMonitoring,
+       WrapKeysetWithMonitoringEncryptDecryptSuccess) {
+  // Populate a primitive set.
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+  auto aead_primitive_set = absl::make_unique<PrimitiveSet<Aead>>(kAnnotations);
+  ASSERT_THAT(aead_primitive_set
+                  ->AddPrimitive(absl::make_unique<DummyAead>("aead0"),
+                                 keyset_info.key_info(0))
+                  .status(),
+              IsOk());
+  ASSERT_THAT(aead_primitive_set
+                  ->AddPrimitive(absl::make_unique<DummyAead>("aead1"),
+                                 keyset_info.key_info(1))
+                  .status(),
+              IsOk());
+  // Set the last as primary.
+  util::StatusOr<PrimitiveSet<Aead>::Entry<Aead>*> last =
+      aead_primitive_set->AddPrimitive(absl::make_unique<DummyAead>("aead2"),
+                                       keyset_info.key_info(2));
+  ASSERT_THAT(last.status(), IsOk());
+  ASSERT_THAT(aead_primitive_set->set_primary(*last), IsOk());
+  // Record the ID of the primary key.
+  const uint32_t kPrimaryKeyId = keyset_info.key_info(2).key_id();
+
+  util::StatusOr<std::unique_ptr<Aead>> aead =
+      AeadWrapper().Wrap(std::move(aead_primitive_set));
+  ASSERT_THAT(aead.status(), IsOk());
+
+  constexpr absl::string_view kPlaintext = "This is some plaintext!";
+  constexpr absl::string_view kAdditionalData = "Some additional data!";
+  EXPECT_CALL(*encryption_monitoring_client_ptr_,
+              Log(kPrimaryKeyId, kPlaintext.size()));
+  util::StatusOr<std::string> ciphertext =
+      (*aead)->Encrypt(kPlaintext, kAdditionalData);
+  ASSERT_THAT(ciphertext.status(), IsOk());
+
+  // In the log expect the size of the ciphertext without the non-raw prefix.
+  auto raw_ciphertext =
+      absl::string_view(*ciphertext).substr(CryptoFormat::kNonRawPrefixSize);
+  EXPECT_CALL(*decryption_monitoring_client_ptr_,
+              Log(kPrimaryKeyId, raw_ciphertext.size()));
+  EXPECT_THAT((*aead)->Decrypt(*ciphertext, kAdditionalData).status(), IsOk());
+}
+
+// Test that monitoring logs encryption and decryption failures correctly.
+TEST_F(AeadSetWrapperTestWithMonitoring,
+       WrapKeysetWithMonitoringEncryptDecryptFailures) {
+  // Populate a primitive set.
+  KeysetInfo keyset_info = CreateTestKeysetInfo();
+
+  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
+      {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
+
+  auto aead_primitive_set = absl::make_unique<PrimitiveSet<Aead>>(kAnnotations);
+
+  // Assume encryption and decryption always fail.
+  auto mock_aead = absl::make_unique<MockAead>();
+  constexpr absl::string_view kPlaintext = "A plaintext!!";
+  constexpr absl::string_view kCiphertext = "A ciphertext!";
+  constexpr absl::string_view kAdditionalData = "Some additional data!";
+  ON_CALL(*mock_aead, Encrypt(kPlaintext, kAdditionalData))
+      .WillByDefault(Return(util::Status(absl::StatusCode::kInternal,
+                                         "Oh no encryption failed :(!")));
+  ON_CALL(*mock_aead, Decrypt(kCiphertext, kAdditionalData))
+      .WillByDefault(Return(util::Status(absl::StatusCode::kInternal,
+                                         "Oh no decryption failed :(!")));
+
+  util::StatusOr<PrimitiveSet<Aead>::Entry<Aead>*> primary =
+      aead_primitive_set->AddPrimitive(std::move(mock_aead),
+                                       keyset_info.key_info(2));
+  ASSERT_THAT(primary.status(), IsOk());
+  // Set the only primitive as primary.
+  ASSERT_THAT(aead_primitive_set->set_primary(*primary), IsOk());
+
+  util::StatusOr<std::unique_ptr<Aead>> aead =
+      AeadWrapper().Wrap(std::move(aead_primitive_set));
+  ASSERT_THAT(aead.status(), IsOk());
+
+  // Expect encryption failure gets logged.
+  EXPECT_CALL(*encryption_monitoring_client_ptr_, LogFailure());
+  util::StatusOr<std::string> ciphertext =
+      (*aead)->Encrypt(kPlaintext, kAdditionalData);
+  EXPECT_THAT(ciphertext.status(), Not(IsOk()));
+
+  // We must prepend the identifier to the ciphertext to make sure our mock gets
+  // called.
+  util::StatusOr<std::string> key_identifier =
+      CryptoFormat::GetOutputPrefix(keyset_info.key_info(2));
+  ASSERT_THAT(key_identifier.status(), IsOk());
+  std::string ciphertext_with_key_id =
+      absl::StrCat(*key_identifier, kCiphertext);
+
+  // Expect decryption failure gets logged.
+  EXPECT_CALL(*decryption_monitoring_client_ptr_, LogFailure());
+  EXPECT_THAT(
+      (*aead)->Decrypt(ciphertext_with_key_id, kAdditionalData).status(),
+      Not(IsOk()));
+}
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto