Add APIs to KeysetHandle to hold annotations for a Keyset.

Annotations are used as contextual information for a Keyset that the application can pass down to Tink, and are used for monitoring.

#tinkApiChange

PiperOrigin-RevId: 438826607

2 files changed, 25 insertions, 288 deletions

diff --git a/cc/core/keyset_handle.cc b/cc/core/keyset_handle.cc
index 3b75cfbf6..9efb91da4 100644
--- a/cc/core/keyset_handle.cc
+++ b/cc/core/keyset_handle.cc
@@ -19,7 +19,6 @@
 #include <string>
 #include <utility>
 
-#include "absl/container/flat_hash_map.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
@@ -85,22 +84,18 @@ util::Status ValidateNoSecret(const Keyset& keyset) {
 
 }  // anonymous namespace
 
+// static
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::Read(
-    std::unique_ptr<KeysetReader> reader, const Aead& master_key_aead,
-    const absl::flat_hash_map<std::string, std::string>&
-        monitoring_annotations) {
-  return ReadWithAssociatedData(std::move(reader), master_key_aead,
-                                /*associated_data=*/"", monitoring_annotations);
+    std::unique_ptr<KeysetReader> reader, const Aead& master_key_aead) {
+  return ReadWithAssociatedData(std::move(reader), master_key_aead, "");
 }
 
+// static
 util::StatusOr<std::unique_ptr<KeysetHandle>>
-KeysetHandle::ReadWithAssociatedData(
-    std::unique_ptr<KeysetReader> reader, const Aead& master_key_aead,
-    absl::string_view associated_data,
-    const absl::flat_hash_map<std::string, std::string>&
-        monitoring_annotations) {
-  util::StatusOr<std::unique_ptr<EncryptedKeyset>> enc_keyset_result =
-      reader->ReadEncrypted();
+KeysetHandle::ReadWithAssociatedData(std::unique_ptr<KeysetReader> reader,
+                                     const Aead& master_key_aead,
+                                     absl::string_view associated_data) {
+  auto enc_keyset_result = reader->ReadEncrypted();
   if (!enc_keyset_result.ok()) {
     return ToStatusF(absl::StatusCode::kInvalidArgument,
                      "Error reading encrypted keyset data: %s",
@@ -114,25 +109,23 @@ KeysetHandle::ReadWithAssociatedData(
                      "Error decrypting encrypted keyset: %s",
                      keyset_result.status().message());
   }
-  return absl::WrapUnique(
-      new KeysetHandle(*std::move(keyset_result), monitoring_annotations));
+
+  std::unique_ptr<KeysetHandle> handle(
+      new KeysetHandle(std::move(keyset_result.value())));
+  return std::move(handle);
 }
 
+// static
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::ReadNoSecret(
-    const std::string& serialized_keyset,
-    const absl::flat_hash_map<std::string, std::string>&
-        monitoring_annotations) {
+    const std::string& serialized_keyset) {
   Keyset keyset;
   if (!keyset.ParseFromString(serialized_keyset)) {
     return util::Status(absl::StatusCode::kInvalidArgument,
                         "Could not parse the input string as a Keyset-proto.");
   }
   util::Status validation = ValidateNoSecret(keyset);
-  if (!validation.ok()) {
-    return validation;
-  }
-  return absl::WrapUnique(
-      new KeysetHandle(std::move(keyset), monitoring_annotations));
+  if (!validation.ok()) return validation;
+  return absl::WrapUnique(new KeysetHandle(std::move(keyset)));
 }
 
 util::Status KeysetHandle::Write(KeysetWriter* writer,
@@ -168,18 +161,15 @@ util::Status KeysetHandle::WriteNoSecret(KeysetWriter* writer) const {
   return writer->Write(get_keyset());
 }
 
+// static
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::GenerateNew(
-    const KeyTemplate& key_template,
-    const absl::flat_hash_map<std::string, std::string>&
-        monitoring_annotations) {
+    const KeyTemplate& key_template) {
   Keyset keyset;
-  util::StatusOr<uint32_t> const result =
-      AddToKeyset(key_template, /*as_primary=*/true, &keyset);
+  auto result = AddToKeyset(key_template, /*as_primary=*/true, &keyset);
   if (!result.ok()) {
     return result.status();
   }
-  return absl::WrapUnique(
-      new KeysetHandle(std::move(keyset), monitoring_annotations));
+  return absl::WrapUnique<KeysetHandle>(new KeysetHandle(std::move(keyset)));
 }
 
 util::StatusOr<std::unique_ptr<Keyset::Key>> ExtractPublicKey(
diff --git a/cc/core/keyset_handle_test.cc b/cc/core/keyset_handle_test.cc
index 0e4abb81c..1381ec651 100644
--- a/cc/core/keyset_handle_test.cc
+++ b/cc/core/keyset_handle_test.cc
@@ -22,10 +22,8 @@
 #include <string>
 #include <utility>
 
-#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "absl/status/status.h"
-#include "absl/strings/string_view.h"
 #include "tink/aead/aead_key_templates.h"
 #include "tink/aead/aead_wrapper.h"
 #include "tink/aead/aes_gcm_key_manager.h"
@@ -37,12 +35,9 @@
 #include "tink/core/key_manager_impl.h"
 #include "tink/json_keyset_reader.h"
 #include "tink/json_keyset_writer.h"
-#include "tink/primitive_set.h"
-#include "tink/primitive_wrapper.h"
 #include "tink/signature/ecdsa_sign_key_manager.h"
 #include "tink/signature/signature_key_templates.h"
 #include "tink/util/protobuf_helper.h"
-#include "tink/util/status.h"
 #include "tink/util/test_keyset_handle.h"
 #include "tink/util/test_matchers.h"
 #include "tink/util/test_util.h"
@@ -66,7 +61,6 @@ using google::crypto::tink::Keyset;
 using google::crypto::tink::KeyStatusType;
 using google::crypto::tink::KeyTemplate;
 using google::crypto::tink::OutputPrefixType;
-using ::testing::_;
 using ::testing::Not;
 
 namespace {
@@ -79,92 +73,6 @@ class KeysetHandleTest : public ::testing::Test {
   }
 };
 
-// Dummy key factory that is required to create a key manager.
-class DummyAeadKeyFactory : public KeyFactory {
- public:
-  explicit DummyAeadKeyFactory(absl::string_view key_type)
-      : key_type_(key_type) {}
-
-  util::StatusOr<std::unique_ptr<portable_proto::MessageLite>> NewKey(
-      const portable_proto::MessageLite& key_format) const override {
-    return util::Status(absl::StatusCode::kUnimplemented, "Unimplemented");
-  }
-
-  util::StatusOr<std::unique_ptr<portable_proto::MessageLite>> NewKey(
-      absl::string_view serialized_key_format) const override {
-    return util::Status(absl::StatusCode::kUnimplemented, "Unimplemented");
-  }
-
-  util::StatusOr<std::unique_ptr<KeyData>> NewKeyData(
-      absl::string_view serialized_key_format) const override {
-    auto key_data = absl::make_unique<KeyData>();
-    key_data->set_type_url(key_type_);
-    std::string serialized_key_format_str(serialized_key_format);
-    key_data->set_value(serialized_key_format_str);
-    return std::move(key_data);
-  }
-
- private:
-  const std::string key_type_;
-};
-
-// Fake Aead key manager for testing.
-class FakeAeadKeyManager : public KeyManager<Aead> {
- public:
-  explicit FakeAeadKeyManager(absl::string_view key_type)
-      : key_type_(key_type), key_factory_(key_type) {}
-
-  util::StatusOr<std::unique_ptr<Aead>> GetPrimitive(
-      const KeyData& key) const override {
-    return {absl::make_unique<DummyAead>(key_type_)};
-  }
-
-  util::StatusOr<std::unique_ptr<Aead>> GetPrimitive(
-      const portable_proto::MessageLite& key) const override {
-    return util::Status(absl::StatusCode::kUnknown,
-                        "DummyAeadKeyFactory cannot construct an aead");
-  }
-
-  uint32_t get_version() const override { return 0; }
-  const std::string& get_key_type() const override { return key_type_; }
-  const KeyFactory& get_key_factory() const override { return key_factory_; }
-
- private:
-  const std::string key_type_;
-  const DummyAeadKeyFactory key_factory_;
-};
-
-class MockAeadPrimitiveWrapper : public PrimitiveWrapper<Aead, Aead> {
- public:
-  MOCK_METHOD(util::StatusOr<std::unique_ptr<Aead>>, Wrap,
-              (std::unique_ptr<PrimitiveSet<Aead>> primitive_set),
-              (const override));
-};
-
-// Generates a keyset for testing.
-Keyset GetTestKeyset() {
-  Keyset keyset;
-  Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
-             KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
-            KeyData::SYMMETRIC, &keyset);
-  keyset.set_primary_key_id(42);
-  return keyset;
-}
-
-// Generates a public keyset for testing.
-Keyset GetPublicTestKeyset() {
-  Keyset keyset;
-  Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
-             KeyData::ASYMMETRIC_PUBLIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
-            KeyData::REMOTE, &keyset);
-  keyset.set_primary_key_id(42);
-  return keyset;
-}
-
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetBinary) {
   Keyset keyset;
   Keyset::Key key;
@@ -232,54 +140,6 @@ TEST_F(KeysetHandleTest, ReadEncryptedKeysetBinary) {
   }
 }
 
-// Check that the generated keyset handle correctly propagates annotations.
-TEST_F(KeysetHandleTest, ReadEncryptedWithAnnotations) {
-  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
-      {"key1", "value1"}, {"key2", "value2"}};
-  Keyset keyset = GetTestKeyset();
-  DummyAead aead("dummy aead 42");
-  std::string keyset_ciphertext =
-      *aead.Encrypt(keyset.SerializeAsString(), /*associated_data=*/"");
-  EncryptedKeyset encrypted_keyset;
-  encrypted_keyset.set_encrypted_keyset(keyset_ciphertext);
-  util::StatusOr<std::unique_ptr<KeysetReader>> reader =
-      BinaryKeysetReader::New(encrypted_keyset.SerializeAsString());
-  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
-      KeysetHandle::Read(*std::move(reader), aead, kAnnotations);
-  ASSERT_THAT(keyset_handle.status(), IsOk());
-
-  // In order to validate annotations are set correctly, we need acceess to the
-  // generated primitive set, which is populated by KeysetWrapperImpl and passed
-  // to the primitive wrapper. We thus register a mock primitive wrapper for
-  // Aead so that we can copy the annotations and later check them.
-  auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
-  absl::flat_hash_map<std::string, std::string> generated_annotations;
-  EXPECT_CALL(*primitive_wrapper, Wrap(_))
-      .WillOnce(
-          [&generated_annotations](
-              std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
-            generated_annotations = generated_primitive_set->get_annotations();
-            std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
-            return aead;
-          });
-  Registry::Reset();
-  ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-
-  ASSERT_THAT((*keyset_handle)->GetPrimitive<Aead>().status(), IsOk());
-  EXPECT_EQ(generated_annotations, kAnnotations);
-  // This is needed to cleanup mocks.
-  Registry::Reset();
-}
-
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetJson) {
   Keyset keyset;
   Keyset::Key key;
@@ -427,52 +287,6 @@ TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataGoodKeyset) {
             TestKeysetHandle::GetKeyset(*handle).SerializeAsString());
 }
 
-// Check that the generated keyset handle correctly propagates annotations.
-TEST_F(KeysetHandleTest, ReadEncryptedWithAssociatedDataAndAnnotations) {
-  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
-      {"key1", "value1"}, {"key2", "value2"}};
-  constexpr absl::string_view kAssociatedData = "some associated data";
-  Keyset keyset = GetTestKeyset();
-  DummyAead aead("dummy aead 42");
-  std::string keyset_ciphertext =
-      *aead.Encrypt(keyset.SerializeAsString(), kAssociatedData);
-  EncryptedKeyset encrypted_keyset;
-  encrypted_keyset.set_encrypted_keyset(keyset_ciphertext);
-  util::StatusOr<std::unique_ptr<KeysetReader>> reader =
-      BinaryKeysetReader::New(encrypted_keyset.SerializeAsString());
-  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
-      KeysetHandle::ReadWithAssociatedData(*std::move(reader), aead,
-                                           kAssociatedData, kAnnotations);
-  ASSERT_THAT(keyset_handle.status(), IsOk());
-
-  auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
-  absl::flat_hash_map<std::string, std::string> generated_annotations;
-  EXPECT_CALL(*primitive_wrapper, Wrap(_))
-      .WillOnce(
-          [&generated_annotations](
-              std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
-            generated_annotations = generated_primitive_set->get_annotations();
-            std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
-            return aead;
-          });
-  Registry::Reset();
-  ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-
-  ASSERT_THAT((*keyset_handle)->GetPrimitive<Aead>().status(), IsOk());
-  EXPECT_EQ(generated_annotations, kAnnotations);
-  // This is needed to cleanup mocks.
-  Registry::Reset();
-}
-
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataWrongAad) {
   Keyset keyset;
   Keyset::Key key;
@@ -558,10 +372,10 @@ TEST_F(KeysetHandleTest, WriteEncryptedKeysetWithAssociatedData) {
 
 TEST_F(KeysetHandleTest, GenerateNewKeysetHandle) {
   const google::crypto::tink::KeyTemplate* key_templates[] = {
-      &AeadKeyTemplates::Aes128Gcm(),
-      &AeadKeyTemplates::Aes256Gcm(),
-      &AeadKeyTemplates::Aes128CtrHmacSha256(),
-      &AeadKeyTemplates::Aes256CtrHmacSha256(),
+    &AeadKeyTemplates::Aes128Gcm(),
+    &AeadKeyTemplates::Aes256Gcm(),
+    &AeadKeyTemplates::Aes128CtrHmacSha256(),
+    &AeadKeyTemplates::Aes256CtrHmacSha256(),
   };
   for (auto templ : key_templates) {
     auto handle_result = KeysetHandle::GenerateNew(*templ);
@@ -571,39 +385,6 @@ TEST_F(KeysetHandleTest, GenerateNewKeysetHandle) {
   }
 }
 
-TEST_F(KeysetHandleTest, GenerateNewWithAnnotations) {
-  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
-      {"key1", "value1"}, {"key2", "value2"}};
-
-  // The template used doesn't make any different w.r.t. annotations.
-  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
-      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), kAnnotations);
-  ASSERT_THAT(keyset_handle.status(), IsOk());
-  auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
-  absl::flat_hash_map<std::string, std::string> generated_annotations;
-  EXPECT_CALL(*primitive_wrapper, Wrap(_))
-      .WillOnce(
-          [&generated_annotations](
-              std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
-            generated_annotations = generated_primitive_set->get_annotations();
-            std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
-            return aead;
-          });
-  Registry::Reset();
-  ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>(
-                      "type.googleapis.com/google.crypto.tink.AesGcmKey"),
-                  true),
-              IsOk());
-
-  EXPECT_THAT((*keyset_handle)->GetPrimitive<Aead>().status(), IsOk());
-  EXPECT_EQ(generated_annotations, kAnnotations);
-  // This is needed to cleanup mocks.
-  Registry::Reset();
-}
-
 TEST_F(KeysetHandleTest, GenerateNewKeysetHandleErrors) {
   KeyTemplate templ;
   templ.set_type_url("type.googleapis.com/some.unknown.KeyType");
@@ -683,6 +464,7 @@ TEST_F(KeysetHandleTest, GetPublicKeysetHandle) {
   }
 }
 
+
 TEST_F(KeysetHandleTest, GetPublicKeysetHandleErrors) {
   { // A keyset with a single key.
     auto handle_result = KeysetHandle::GenerateNew(
@@ -821,41 +603,6 @@ TEST_F(KeysetHandleTest, ReadNoSecret) {
   ASSERT_EQ(result.key(1).key_id(), keyset.key(1).key_id());
 }
 
-TEST_F(KeysetHandleTest, ReadNoSecretWithAnnotations) {
-  const absl::flat_hash_map<std::string, std::string> kAnnotations = {
-      {"key1", "value1"}, {"key2", "value2"}};
-  Keyset keyset = GetPublicTestKeyset();
-  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
-      KeysetHandle::ReadNoSecret(keyset.SerializeAsString(), kAnnotations);
-  ASSERT_THAT(keyset_handle.status(), IsOk());
-  auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
-  absl::flat_hash_map<std::string, std::string> generated_annotations;
-  EXPECT_CALL(*primitive_wrapper, Wrap(_))
-      .WillOnce(
-          [&generated_annotations](
-              std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
-            generated_annotations = generated_primitive_set->get_annotations();
-            std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
-            return aead;
-          });
-  Registry::Reset();
-  ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
-                  /*new_key_allowed=*/true),
-              IsOk());
-
-  EXPECT_THAT((*keyset_handle)->GetPrimitive<Aead>().status(), IsOk());
-  EXPECT_EQ(generated_annotations, kAnnotations);
-  // This is needed to cleanup mocks.
-  Registry::Reset();
-}
-
 TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeUnknown) {
   Keyset keyset;
   Keyset::Key key;