Replace unique_ptr with shared_ptr in KeysetHandleBuilder::Entry API.

Also, return shared_ptr in KeysetHandle::Entry::GetKey().

PiperOrigin-RevId: 523772662

6 files changed, 226 insertions, 88 deletions

diff --git a/cc/core/cleartext_keyset_handle.cc b/cc/core/cleartext_keyset_handle.cc
index 989f95e35..9e000ff18 100644
--- a/cc/core/cleartext_keyset_handle.cc
+++ b/cc/core/cleartext_keyset_handle.cc
@@ -20,6 +20,7 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
@@ -41,14 +42,23 @@ util::StatusOr<std::unique_ptr<KeysetHandle>> CleartextKeysetHandle::Read(
     std::unique_ptr<KeysetReader> reader,
     const absl::flat_hash_map<std::string, std::string>&
         monitoring_annotations) {
-  auto keyset_result = reader->Read();
+  util::StatusOr<std::unique_ptr<Keyset>> keyset_result = reader->Read();
   if (!keyset_result.ok()) {
     return ToStatusF(absl::StatusCode::kInvalidArgument,
                      "Error reading keyset data: %s",
                      keyset_result.status().message());
   }
+  util::StatusOr<std::vector<std::shared_ptr<const KeysetHandle::Entry>>>
+      entries = KeysetHandle::GetEntriesFromKeyset(**keyset_result);
+  if (!entries.ok()) {
+    return entries.status();
+  }
+  if (entries->size() != (*keyset_result)->key_size()) {
+    return util::Status(absl::StatusCode::kInternal,
+                        "Error converting keyset proto into key entries.");
+  }
   std::unique_ptr<KeysetHandle> handle(new KeysetHandle(
-      std::move(keyset_result.value()), monitoring_annotations));
+      std::move(keyset_result.value()), *entries, monitoring_annotations));
   return std::move(handle);
 }
 
diff --git a/cc/core/cleartext_keyset_handle_test.cc b/cc/core/cleartext_keyset_handle_test.cc
index 428204018..0ff416c6a 100644
--- a/cc/core/cleartext_keyset_handle_test.cc
+++ b/cc/core/cleartext_keyset_handle_test.cc
@@ -48,9 +48,9 @@ class CleartextKeysetHandleTest : public ::testing::Test {
 TEST_F(CleartextKeysetHandleTest, testRead) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   {  // Reader that reads a valid keyset.
@@ -75,9 +75,9 @@ TEST_F(CleartextKeysetHandleTest, testRead) {
 TEST_F(CleartextKeysetHandleTest, testWrite) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
 
diff --git a/cc/core/keyset_handle.cc b/cc/core/keyset_handle.cc
index 0ee17d7fb..bab5bfe99 100644
--- a/cc/core/keyset_handle.cc
+++ b/cc/core/keyset_handle.cc
@@ -15,10 +15,11 @@
 ///////////////////////////////////////////////////////////////////////////////
 #include "tink/keyset_handle.h"
 
-#include <iostream>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
@@ -179,29 +180,57 @@ KeysetHandle::Entry KeysetHandle::operator[](int index) const {
   CHECK(index >= 0 && index < size())
       << "Invalid index " << index << " for keyset of size " << size();
 
+  if (!entries_.empty() && entries_.size() > index) {
+    return *entries_[index];
+  }
+  // Since `entries_` has not been populated, the entry must be created on
+  // demand from the key proto entry at `index` in `keyset_`. This special
+  // case will no longer be necessary after `keyset_` has been removed from the
+  // `KeysetHandle` class.
+  //
+  // TODO(b/277792846): Remove after transition to rely solely on
+  // `KeysetHandle::Entry`.
+  return CreateEntryAt(index);
+}
+
+KeysetHandle::Entry KeysetHandle::CreateEntryAt(int index) const {
+  CHECK(index >= 0 && index < size())
+      << "Invalid index " << index << " for keyset of size " << size();
+
   util::Status validation = ValidateAt(index);
   CHECK_OK(validation);
 
-  const Keyset::Key& proto_key = get_keyset().key(index);
-  int id = proto_key.key_id();
+  Keyset keyset = get_keyset();
+  util::StatusOr<Entry> entry =
+      CreateEntry(keyset.key(index), keyset.primary_key_id());
+  // Status should be OK since this keyset handle has been validated.
+  CHECK_OK(entry.status());
+  return *entry;
+}
 
+util::StatusOr<KeysetHandle::Entry> KeysetHandle::CreateEntry(
+    const Keyset::Key& proto_key, uint32_t primary_key_id) {
   util::StatusOr<internal::ProtoKeySerialization> serialization =
       ToProtoKeySerialization(proto_key);
-  // Status should be OK since this keyset handle has been validated.
-  CHECK_OK(serialization.status());
+  if (!serialization.ok()) {
+    return serialization.status();
+  }
 
-  util::StatusOr<std::unique_ptr<Key>> key =
+  util::StatusOr<std::shared_ptr<const Key>> key =
       internal::MutableSerializationRegistry::GlobalInstance()
           .ParseKeyWithLegacyFallback(*serialization);
-  CHECK_OK(key.status());
+  if (!key.ok()) {
+    return key.status();
+  }
 
   util::StatusOr<KeyStatus> key_status =
       internal::FromKeyStatusType(proto_key.status());
-  // Status should be OK since this keyset handle has been validated.
-  CHECK_OK(key_status.status());
+  if (!key_status.ok()) {
+    return key_status.status();
+  }
 
-  return Entry(*std::move(key), *key_status, id,
-               id == get_keyset().primary_key_id());
+  return Entry(*std::move(key), *key_status, proto_key.key_id(),
+               proto_key.key_id() == primary_key_id);
 }
 
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::Read(
@@ -233,8 +262,17 @@ KeysetHandle::ReadWithAssociatedData(
                      "Error decrypting encrypted keyset: %s",
                      keyset_result.status().message());
   }
-  return absl::WrapUnique(
-      new KeysetHandle(*std::move(keyset_result), monitoring_annotations));
+  util::StatusOr<std::vector<std::shared_ptr<const Entry>>> entries =
+      GetEntriesFromKeyset(**keyset_result);
+  if (!entries.ok()) {
+    return entries.status();
+  }
+  if (entries->size() != (*keyset_result)->key_size()) {
+    return util::Status(absl::StatusCode::kInternal,
+                        "Error converting keyset proto into key entries.");
+  }
+  return absl::WrapUnique(new KeysetHandle(*std::move(keyset_result), *entries,
+                                           monitoring_annotations));
 }
 
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::ReadNoSecret(
@@ -250,8 +288,17 @@ util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::ReadNoSecret(
   if (!validation.ok()) {
     return validation;
   }
+  util::StatusOr<std::vector<std::shared_ptr<const Entry>>> entries =
+      GetEntriesFromKeyset(keyset);
+  if (!entries.ok()) {
+    return entries.status();
+  }
+  if (entries->size() != keyset.key_size()) {
+    return util::Status(absl::StatusCode::kInternal,
+                        "Error converting keyset proto into key entries.");
+  }
   return absl::WrapUnique(
-      new KeysetHandle(std::move(keyset), monitoring_annotations));
+      new KeysetHandle(std::move(keyset), *entries, monitoring_annotations));
 }
 
 util::Status KeysetHandle::Write(KeysetWriter* writer,
@@ -325,8 +372,17 @@ KeysetHandle::GetPublicKeysetHandle() const {
     public_keyset->add_key()->Swap(public_key_result.value().get());
   }
   public_keyset->set_primary_key_id(get_keyset().primary_key_id());
+  util::StatusOr<std::vector<std::shared_ptr<const Entry>>> entries =
+      GetEntriesFromKeyset(*public_keyset);
+  if (!entries.ok()) {
+    return entries.status();
+  }
+  if (entries->size() != public_keyset->key_size()) {
+    return util::Status(absl::StatusCode::kInternal,
+                        "Error converting keyset proto into key entries.");
+  }
   std::unique_ptr<KeysetHandle> handle(
-      new KeysetHandle(std::move(public_keyset)));
+      new KeysetHandle(std::move(public_keyset), *entries));
   return std::move(handle);
 }
 
@@ -355,19 +411,36 @@ crypto::tink::util::StatusOr<uint32_t> KeysetHandle::AddToKeyset(
 
 crypto::tink::util::StatusOr<uint32_t> KeysetHandle::AddKey(
     const google::crypto::tink::KeyTemplate& key_template, bool as_primary) {
-  return AddToKeyset(key_template, as_primary, &keyset_);
+  util::StatusOr<uint32_t> id = AddToKeyset(key_template, as_primary, &keyset_);
+  if (!id.ok()) {
+    return id.status();
+  }
+  util::StatusOr<const Entry> entry = CreateEntry(
+      keyset_.key(keyset_.key_size() - 1), keyset_.primary_key_id());
+  if (!entry.ok()) {
+    return entry.status();
+  }
+  entries_.push_back(std::make_shared<const Entry>(*entry));
+  return *id;
 }
 
 KeysetInfo KeysetHandle::GetKeysetInfo() const {
   return KeysetInfoFromKeyset(get_keyset());
 }
 
-KeysetHandle::KeysetHandle(Keyset keyset) : keyset_(std::move(keyset)) {}
-
-KeysetHandle::KeysetHandle(std::unique_ptr<Keyset> keyset)
-    : keyset_(std::move(*keyset)) {}
-
-const Keyset& KeysetHandle::get_keyset() const { return keyset_; }
+util::StatusOr<std::vector<std::shared_ptr<const KeysetHandle::Entry>>>
+KeysetHandle::GetEntriesFromKeyset(const Keyset& keyset) {
+  std::vector<std::shared_ptr<const Entry>> entries;
+  for (const Keyset::Key& key : keyset.key()) {
+    util::StatusOr<const Entry> entry =
+        CreateEntry(key, keyset.primary_key_id());
+    if (!entry.ok()) {
+      return entry.status();
+    }
+    entries.push_back(std::make_shared<const Entry>(*entry));
+  }
+  return entries;
+}
 
 }  // namespace tink
 }  // namespace crypto
diff --git a/cc/core/keyset_handle_builder.cc b/cc/core/keyset_handle_builder.cc
index 2dffe6843..223321bf0 100644
--- a/cc/core/keyset_handle_builder.cc
+++ b/cc/core/keyset_handle_builder.cc
@@ -67,7 +67,7 @@ KeysetHandleBuilder::KeysetHandleBuilder(const KeysetHandle& handle) {
 }
 
 KeysetHandleBuilder::Entry KeysetHandleBuilder::Entry::CreateFromKey(
-    std::unique_ptr<Key> key, KeyStatus status, bool is_primary) {
+    std::shared_ptr<const Key> key, KeyStatus status, bool is_primary) {
   absl::optional<int> id_requirement = key->GetIdRequirement();
   auto imported_entry = absl::make_unique<internal::KeyEntry>(std::move(key));
   KeysetHandleBuilder::Entry entry(std::move(imported_entry));
@@ -76,8 +76,8 @@ KeysetHandleBuilder::Entry KeysetHandleBuilder::Entry::CreateFromKey(
 }
 
 KeysetHandleBuilder::Entry KeysetHandleBuilder::Entry::CreateFromParams(
-    std::unique_ptr<Parameters> parameters, KeyStatus status, bool is_primary,
-    absl::optional<int> id) {
+    std::shared_ptr<const Parameters> parameters, KeyStatus status,
+    bool is_primary, absl::optional<int> id) {
   auto generated_entry =
       absl::make_unique<internal::ParametersEntry>(std::move(parameters));
   KeysetHandleBuilder::Entry entry(std::move(generated_entry));
@@ -189,7 +189,9 @@ util::StatusOr<KeysetHandle> KeysetHandleBuilder::Build() {
                         "No primary set in this keyset.");
   }
   keyset.set_primary_key_id(*primary_id);
-  return KeysetHandle(keyset);
+  util::StatusOr<std::vector<std::shared_ptr<const KeysetHandle::Entry>>>
+      entries = KeysetHandle::GetEntriesFromKeyset(keyset);
+  return KeysetHandle(keyset, *std::move(entries));
 }
 
 }  // namespace tink
diff --git a/cc/core/keyset_handle_builder_test.cc b/cc/core/keyset_handle_builder_test.cc
index 9f3806370..3d2959334 100644
--- a/cc/core/keyset_handle_builder_test.cc
+++ b/cc/core/keyset_handle_builder_test.cc
@@ -51,6 +51,7 @@ namespace {
 
 using ::crypto::tink::test::AddTinkKey;
 using ::crypto::tink::test::IsOk;
+using ::crypto::tink::test::IsOkAndHolds;
 using ::crypto::tink::test::StatusIs;
 using ::google::crypto::tink::AesCmacParams;
 using ::google::crypto::tink::KeyData;
@@ -101,7 +102,7 @@ TEST_F(KeysetHandleBuilderTest, BuildWithSingleKey) {
   EXPECT_THAT((*handle)[0].GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT((*handle)[0].GetId(), Eq(123));
   EXPECT_THAT((*handle)[0].IsPrimary(), IsTrue());
-  EXPECT_THAT((*handle)[0].GetKey().GetParameters().HasIdRequirement(),
+  EXPECT_THAT((*handle)[0].GetKey()->GetParameters().HasIdRequirement(),
               IsTrue());
 }
 
@@ -137,19 +138,19 @@ TEST_F(KeysetHandleBuilderTest, BuildWithMultipleKeys) {
   EXPECT_THAT((*handle)[0].GetStatus(), Eq(KeyStatus::kDestroyed));
   EXPECT_THAT((*handle)[0].GetId(), Eq(123));
   EXPECT_THAT((*handle)[0].IsPrimary(), IsFalse());
-  EXPECT_THAT((*handle)[0].GetKey().GetParameters().HasIdRequirement(),
+  EXPECT_THAT((*handle)[0].GetKey()->GetParameters().HasIdRequirement(),
               IsTrue());
 
   EXPECT_THAT((*handle)[1].GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT((*handle)[1].GetId(), Eq(456));
   EXPECT_THAT((*handle)[1].IsPrimary(), IsTrue());
-  EXPECT_THAT((*handle)[1].GetKey().GetParameters().HasIdRequirement(),
+  EXPECT_THAT((*handle)[1].GetKey()->GetParameters().HasIdRequirement(),
               IsTrue());
 
   EXPECT_THAT((*handle)[2].GetStatus(), Eq(KeyStatus::kDisabled));
   EXPECT_THAT((*handle)[2].GetId(), Eq(789));
   EXPECT_THAT((*handle)[2].IsPrimary(), IsFalse());
-  EXPECT_THAT((*handle)[2].GetKey().GetParameters().HasIdRequirement(),
+  EXPECT_THAT((*handle)[2].GetKey()->GetParameters().HasIdRequirement(),
               IsTrue());
 }
 
@@ -188,17 +189,20 @@ TEST_F(KeysetHandleBuilderTest, BuildCopy) {
   EXPECT_THAT((*copy)[0].GetStatus(), Eq(KeyStatus::kDestroyed));
   EXPECT_THAT((*copy)[0].GetId(), Eq(123));
   EXPECT_THAT((*copy)[0].IsPrimary(), IsFalse());
-  EXPECT_THAT((*copy)[0].GetKey().GetParameters().HasIdRequirement(), IsTrue());
+  EXPECT_THAT((*copy)[0].GetKey()->GetParameters().HasIdRequirement(),
+              IsTrue());
 
   EXPECT_THAT((*copy)[1].GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT((*copy)[1].GetId(), Eq(456));
   EXPECT_THAT((*copy)[1].IsPrimary(), IsTrue());
-  EXPECT_THAT((*copy)[1].GetKey().GetParameters().HasIdRequirement(), IsTrue());
+  EXPECT_THAT((*copy)[1].GetKey()->GetParameters().HasIdRequirement(),
+              IsTrue());
 
   EXPECT_THAT((*copy)[2].GetStatus(), Eq(KeyStatus::kDisabled));
   EXPECT_THAT((*copy)[2].GetId(), Eq(789));
   EXPECT_THAT((*copy)[2].IsPrimary(), IsFalse());
-  EXPECT_THAT((*copy)[2].GetKey().GetParameters().HasIdRequirement(), IsTrue());
+  EXPECT_THAT((*copy)[2].GetKey()->GetParameters().HasIdRequirement(),
+              IsTrue());
 }
 
 TEST_F(KeysetHandleBuilderTest, IsPrimary) {
@@ -387,7 +391,7 @@ TEST_F(KeysetHandleBuilderTest, RemoveEntry) {
   EXPECT_THAT((*handle1)[0].GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT((*handle1)[0].GetId(), Eq(456));
   EXPECT_THAT((*handle1)[0].IsPrimary(), IsTrue());
-  EXPECT_THAT((*handle1)[0].GetKey().GetParameters().HasIdRequirement(),
+  EXPECT_THAT((*handle1)[0].GetKey()->GetParameters().HasIdRequirement(),
               IsTrue());
 }
 
@@ -779,6 +783,55 @@ TEST_F(KeysetHandleBuilderTest, BuildTwiceFails) {
               StatusIs(absl::StatusCode::kFailedPrecondition));
 }
 
+TEST_F(KeysetHandleBuilderTest, UsePrimitivesFromSplitKeyset) {
+  util::StatusOr<AesCmacParameters> params = AesCmacParameters::Create(
+      /*key_size_in_bytes=*/32, /*cryptographic_tag_size_in_bytes=*/16,
+      AesCmacParameters::Variant::kTink);
+  ASSERT_THAT(params, IsOk());
+
+  util::StatusOr<KeysetHandle> handle =
+      KeysetHandleBuilder()
+          .AddEntry(KeysetHandleBuilder::Entry::CreateFromCopyableParams(
+              *params, KeyStatus::kEnabled, /*is_primary=*/false))
+          .AddEntry(KeysetHandleBuilder::Entry::CreateFromCopyableParams(
+              *params, KeyStatus::kEnabled, /*is_primary=*/true))
+          .Build();
+  ASSERT_THAT(handle, IsOkAndHolds(SizeIs(2)));
+
+  util::StatusOr<KeysetHandle> handle0 =
+      KeysetHandleBuilder()
+          .AddEntry(KeysetHandleBuilder::Entry::CreateFromKey(
+              (*handle)[0].GetKey(), KeyStatus::kEnabled,
+              /*is_primary=*/true))
+          .Build();
+  ASSERT_THAT(handle0, IsOkAndHolds(SizeIs(1)));
+  ASSERT_THAT((*handle)[0].GetId(), Eq((*handle0)[0].GetId()));
+
+  util::StatusOr<KeysetHandle> handle1 =
+      KeysetHandleBuilder()
+          .AddEntry(KeysetHandleBuilder::Entry::CreateFromKey(
+              (*handle)[1].GetKey(), KeyStatus::kEnabled,
+              /*is_primary=*/true))
+          .Build();
+  ASSERT_THAT(handle1, IsOkAndHolds(SizeIs(1)));
+  ASSERT_THAT((*handle)[1].GetId(), Eq((*handle1)[0].GetId()));
+
+  util::StatusOr<std::unique_ptr<Mac>> mac0 = handle0->GetPrimitive<Mac>();
+  ASSERT_THAT(mac0.status(), IsOk());
+  util::StatusOr<std::string> tag0 = (*mac0)->ComputeMac("some input");
+  ASSERT_THAT(tag0.status(), IsOk());
+
+  util::StatusOr<std::unique_ptr<Mac>> mac1 = handle1->GetPrimitive<Mac>();
+  ASSERT_THAT(mac1.status(), IsOk());
+  util::StatusOr<std::string> tag1 = (*mac1)->ComputeMac("some other input");
+  ASSERT_THAT(tag1.status(), IsOk());
+
+  // Use original keyset to verify tags computed from new keysets.
+  util::StatusOr<std::unique_ptr<Mac>> mac = handle->GetPrimitive<Mac>();
+  ASSERT_THAT(mac.status(), IsOk());
+  EXPECT_THAT((*mac)->VerifyMac(*tag0, "some input"), IsOk());
+  EXPECT_THAT((*mac)->VerifyMac(*tag1, "some other input"), IsOk());
+}
 
 }  // namespace
 }  // namespace tink
diff --git a/cc/core/keyset_handle_test.cc b/cc/core/keyset_handle_test.cc
index ab5eeeabd..3098ef87f 100644
--- a/cc/core/keyset_handle_test.cc
+++ b/cc/core/keyset_handle_test.cc
@@ -160,9 +160,9 @@ class MockAeadPrimitiveWrapper : public PrimitiveWrapper<Aead, Aead> {
 Keyset GetTestKeyset() {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   return keyset;
@@ -172,9 +172,9 @@ Keyset GetTestKeyset() {
 Keyset GetPublicTestKeyset() {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::REMOTE, &keyset);
   keyset.set_primary_key_id(42);
   return keyset;
@@ -183,9 +183,9 @@ Keyset GetPublicTestKeyset() {
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetBinary) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -281,11 +281,11 @@ TEST_F(KeysetHandleTest, ReadEncryptedWithAnnotations) {
   ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_other_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
 
@@ -298,9 +298,9 @@ TEST_F(KeysetHandleTest, ReadEncryptedWithAnnotations) {
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetJson) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -381,9 +381,9 @@ TEST_F(KeysetHandleTest, WriteEncryptedKeyset_Json) {
   // Prepare a valid keyset handle
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   auto reader =
@@ -421,9 +421,9 @@ TEST_F(KeysetHandleTest, WriteEncryptedKeyset_Json) {
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataGoodKeyset) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -474,11 +474,11 @@ TEST_F(KeysetHandleTest, ReadEncryptedWithAssociatedDataAndAnnotations) {
   ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_other_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
 
@@ -491,9 +491,9 @@ TEST_F(KeysetHandleTest, ReadEncryptedWithAssociatedDataAndAnnotations) {
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataWrongAad) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   DummyAead aead("dummy aead 42");
@@ -512,9 +512,9 @@ TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataWrongAad) {
 TEST_F(KeysetHandleTest, ReadEncryptedKeysetWithAssociatedDataEmptyAad) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   DummyAead aead("dummy aead 42");
@@ -533,9 +533,9 @@ TEST_F(KeysetHandleTest, WriteEncryptedKeysetWithAssociatedData) {
   // Prepare a valid keyset handle
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   auto reader =
@@ -869,9 +869,9 @@ TEST_F(KeysetHandleTest, Copiable) {
 TEST_F(KeysetHandleTest, ReadNoSecret) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::REMOTE, &keyset);
   keyset.set_primary_key_id(42);
   auto handle_result = KeysetHandle::ReadNoSecret(keyset.SerializeAsString());
@@ -908,11 +908,11 @@ TEST_F(KeysetHandleTest, ReadNoSecretWithAnnotations) {
   ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
   ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>("some other key type"),
+                  absl::make_unique<FakeAeadKeyManager>("some_other_key_type"),
                   /*new_key_allowed=*/true),
               IsOk());
 
@@ -925,7 +925,7 @@ TEST_F(KeysetHandleTest, ReadNoSecretWithAnnotations) {
 TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeUnknown) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::UNKNOWN_KEYMATERIAL, &keyset);
   keyset.set_primary_key_id(42);
   auto result = KeysetHandle::ReadNoSecret(keyset.SerializeAsString());
@@ -935,7 +935,7 @@ TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeUnknown) {
 TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeSymmetric) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
   auto result = KeysetHandle::ReadNoSecret(keyset.SerializeAsString());
@@ -945,7 +945,7 @@ TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeSymmetric) {
 TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeAssymmetricPrivate) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PRIVATE, &keyset);
   keyset.set_primary_key_id(42);
   auto result = KeysetHandle::ReadNoSecret(keyset.SerializeAsString());
@@ -955,13 +955,13 @@ TEST_F(KeysetHandleTest, ReadNoSecretFailForTypeAssymmetricPrivate) {
 TEST_F(KeysetHandleTest, ReadNoSecretFailForHidden) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddTinkKey(absl::StrCat("more key type", i), i, key, KeyStatusType::ENABLED,
                KeyData::ASYMMETRIC_PUBLIC, &keyset);
   }
-  AddRawKey("some other key type", 10, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 10, key, KeyStatusType::ENABLED,
             KeyData::ASYMMETRIC_PRIVATE, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddRawKey(absl::StrCat("more key type", i + 100), i + 100, key,
@@ -982,9 +982,9 @@ TEST_F(KeysetHandleTest, ReadNoSecretFailForInvalidString) {
 TEST_F(KeysetHandleTest, WriteNoSecret) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED,
             KeyData::REMOTE, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -1001,7 +1001,7 @@ TEST_F(KeysetHandleTest, WriteNoSecret) {
 TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeUnknown) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::UNKNOWN_KEYMATERIAL, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -1018,7 +1018,7 @@ TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeUnknown) {
 TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeSymmetric) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -1035,7 +1035,7 @@ TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeSymmetric) {
 TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeAssymmetricPrivate) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PRIVATE, &keyset);
   keyset.set_primary_key_id(42);
 
@@ -1052,13 +1052,13 @@ TEST_F(KeysetHandleTest, WriteNoSecretFailForTypeAssymmetricPrivate) {
 TEST_F(KeysetHandleTest, WriteNoSecretFailForHidden) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddTinkKey(absl::StrCat("more key type", i), i, key, KeyStatusType::ENABLED,
                KeyData::ASYMMETRIC_PUBLIC, &keyset);
   }
-  AddRawKey("some other key type", 10, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 10, key, KeyStatusType::ENABLED,
             KeyData::ASYMMETRIC_PRIVATE, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddRawKey(absl::StrCat("more key type", i + 100), i + 100, key,
@@ -1080,13 +1080,13 @@ TEST_F(KeysetHandleTest, WriteNoSecretFailForHidden) {
 TEST_F(KeysetHandleTest, GetKeysetInfo) {
   Keyset keyset;
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED,
              KeyData::ASYMMETRIC_PUBLIC, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddTinkKey(absl::StrCat("more key type", i), i, key, KeyStatusType::ENABLED,
                KeyData::ASYMMETRIC_PUBLIC, &keyset);
   }
-  AddRawKey("some other key type", 10, key, KeyStatusType::ENABLED,
+  AddRawKey("some_other_key_type", 10, key, KeyStatusType::ENABLED,
             KeyData::ASYMMETRIC_PRIVATE, &keyset);
   for (int i = 0; i < 10; ++i) {
     AddRawKey(absl::StrCat("more key type", i + 100), i + 100, key,
@@ -1125,8 +1125,8 @@ TEST_F(KeysetHandleTest, GetEntryFromSingleKeyKeyset) {
   EXPECT_THAT(entry.GetId(), Eq(11));
   EXPECT_THAT(entry.GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT(entry.IsPrimary(), IsTrue());
-  EXPECT_THAT(entry.GetKey().GetIdRequirement(), Eq(11));
-  EXPECT_THAT(entry.GetKey().GetParameters().HasIdRequirement(), IsTrue());
+  EXPECT_THAT(entry.GetKey()->GetIdRequirement(), Eq(11));
+  EXPECT_THAT(entry.GetKey()->GetParameters().HasIdRequirement(), IsTrue());
 }
 
 TEST_F(KeysetHandleTest, GetEntryFromMultipleKeyKeyset) {
@@ -1150,24 +1150,24 @@ TEST_F(KeysetHandleTest, GetEntryFromMultipleKeyKeyset) {
   EXPECT_THAT(entry0.GetId(), Eq(11));
   EXPECT_THAT(entry0.GetStatus(), Eq(KeyStatus::kDisabled));
   EXPECT_THAT(entry0.IsPrimary(), IsFalse());
-  EXPECT_THAT(entry0.GetKey().GetIdRequirement(), Eq(absl::nullopt));
-  EXPECT_THAT(entry0.GetKey().GetParameters().HasIdRequirement(), IsFalse());
+  EXPECT_THAT(entry0.GetKey()->GetIdRequirement(), Eq(absl::nullopt));
+  EXPECT_THAT(entry0.GetKey()->GetParameters().HasIdRequirement(), IsFalse());
 
   ASSERT_THAT(handle->ValidateAt(1), IsOk());
   KeysetHandle::Entry entry1 = (*handle)[1];
   EXPECT_THAT(entry1.GetId(), Eq(22));
   EXPECT_THAT(entry1.GetStatus(), Eq(KeyStatus::kEnabled));
   EXPECT_THAT(entry1.IsPrimary(), IsTrue());
-  EXPECT_THAT(entry1.GetKey().GetIdRequirement(), Eq(22));
-  EXPECT_THAT(entry1.GetKey().GetParameters().HasIdRequirement(), IsTrue());
+  EXPECT_THAT(entry1.GetKey()->GetIdRequirement(), Eq(22));
+  EXPECT_THAT(entry1.GetKey()->GetParameters().HasIdRequirement(), IsTrue());
 
   ASSERT_THAT(handle->ValidateAt(2), IsOk());
   KeysetHandle::Entry entry2 = (*handle)[2];
   EXPECT_THAT(entry2.GetId(), Eq(33));
   EXPECT_THAT(entry2.GetStatus(), Eq(KeyStatus::kDestroyed));
   EXPECT_THAT(entry2.IsPrimary(), IsFalse());
-  EXPECT_THAT(entry2.GetKey().GetIdRequirement(), Eq(absl::nullopt));
-  EXPECT_THAT(entry2.GetKey().GetParameters().HasIdRequirement(), IsFalse());
+  EXPECT_THAT(entry2.GetKey()->GetIdRequirement(), Eq(absl::nullopt));
+  EXPECT_THAT(entry2.GetKey()->GetParameters().HasIdRequirement(), IsFalse());
 }
 
 TEST_F(KeysetHandleDeathTest, EntryWithIndexOutOfBoundsCrashes) {