Replace unique_ptr with shared_ptr in KeysetHandleBuilder::Entry API.

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

PiperOrigin-RevId: 524090677

12 files changed, 309 insertions, 122 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) {
diff --git a/cc/internal/keyset_handle_builder_entry.h b/cc/internal/keyset_handle_builder_entry.h
index 437147263..a1d0a6c4c 100644
--- a/cc/internal/keyset_handle_builder_entry.h
+++ b/cc/internal/keyset_handle_builder_entry.h
@@ -97,13 +97,13 @@ class KeyEntry : public KeysetHandleBuilderEntry {
   KeyEntry(const KeyEntry& other) = delete;
   KeyEntry& operator=(const KeyEntry& other) = delete;
 
-  explicit KeyEntry(std::unique_ptr<Key> key) : key_(std::move(key)) {}
+  explicit KeyEntry(std::shared_ptr<const Key> key) : key_(std::move(key)) {}
 
   crypto::tink::util::StatusOr<google::crypto::tink::Keyset::Key>
   CreateKeysetKey(int id) override;
 
  private:
-  std::unique_ptr<Key> key_;
+  std::shared_ptr<const Key> key_;
 };
 
 // Internal keyset handle builder entry constructed from a `Parameters` object.
@@ -115,14 +115,14 @@ class ParametersEntry : public KeysetHandleBuilderEntry {
   ParametersEntry(const ParametersEntry& other) = delete;
   ParametersEntry& operator=(const ParametersEntry& other) = delete;
 
-  explicit ParametersEntry(std::unique_ptr<Parameters> parameters)
+  explicit ParametersEntry(std::shared_ptr<const Parameters> parameters)
       : parameters_(std::move(parameters)) {}
 
   crypto::tink::util::StatusOr<google::crypto::tink::Keyset::Key>
   CreateKeysetKey(int id) override;
 
  private:
-  std::unique_ptr<Parameters> parameters_;
+  std::shared_ptr<const Parameters> parameters_;
 };
 
 }  // namespace internal
diff --git a/cc/keyset_handle.h b/cc/keyset_handle.h
index d94c61e8f..62e09ca68 100644
--- a/cc/keyset_handle.h
+++ b/cc/keyset_handle.h
@@ -17,9 +17,11 @@
 #ifndef TINK_KEYSET_HANDLE_H_
 #define TINK_KEYSET_HANDLE_H_
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/base/attributes.h"
 #include "absl/container/flat_hash_map.h"
@@ -53,7 +55,7 @@ class KeysetHandle {
     // May return an internal class in case there is no implementation of the
     // corresponding key class yet.  Returned value only valid for lifetime
     // of entry object.
-    const Key& GetKey() const { return *key_; }
+    std::shared_ptr<const Key> GetKey() const { return key_; }
 
     // Status indicates whether or not a key should still be used.
     KeyStatus GetStatus() const { return status_; }
@@ -70,13 +72,14 @@ class KeysetHandle {
     friend class KeysetHandle;
     friend class KeysetHandleBuilder;
 
-    Entry(std::unique_ptr<Key> key, KeyStatus status, int id, bool is_primary)
+    Entry(std::shared_ptr<const Key> key, KeyStatus status, int id,
+          bool is_primary)
         : key_(std::move(key)),
           status_(status),
           id_(id),
           is_primary_(is_primary) {}
 
-    std::unique_ptr<Key> key_;
+    std::shared_ptr<const Key> key_;
     KeyStatus status_;
     int id_;
     bool is_primary_;
@@ -199,9 +202,19 @@ class KeysetHandle {
   friend class KeysetHandleBuilder;
 
   // Creates a handle that contains the given keyset.
-  explicit KeysetHandle(google::crypto::tink::Keyset keyset);
-  // Creates a handle that contains the given keyset.
-  explicit KeysetHandle(std::unique_ptr<google::crypto::tink::Keyset> keyset);
+  explicit KeysetHandle(google::crypto::tink::Keyset keyset)
+      : keyset_(std::move(keyset)) {}
+  explicit KeysetHandle(std::unique_ptr<google::crypto::tink::Keyset> keyset)
+      : keyset_(std::move(*keyset)) {}
+  // Creates a handle that contains the given `keyset` and `entries`.
+  explicit KeysetHandle(
+      google::crypto::tink::Keyset keyset,
+      const std::vector<std::shared_ptr<const Entry>>& entries)
+      : keyset_(std::move(keyset)), entries_(entries) {}
+  explicit KeysetHandle(
+      std::unique_ptr<google::crypto::tink::Keyset> keyset,
+      const std::vector<std::shared_ptr<const Entry>>& entries)
+      : keyset_(std::move(*keyset)), entries_(entries) {}
   // Creates a handle that contains the given `keyset` and
   // `monitoring_annotations`.
   KeysetHandle(google::crypto::tink::Keyset keyset,
@@ -214,18 +227,43 @@ class KeysetHandle {
                    monitoring_annotations)
       : keyset_(std::move(*keyset)),
         monitoring_annotations_(monitoring_annotations) {}
+  // Creates a handle that contains the given `keyset`, `entries`, and
+  // `monitoring_annotations`.
+  KeysetHandle(google::crypto::tink::Keyset keyset,
+               const std::vector<std::shared_ptr<const Entry>>& entries,
+               const absl::flat_hash_map<std::string, std::string>&
+                   monitoring_annotations)
+      : keyset_(std::move(keyset)),
+        entries_(entries),
+        monitoring_annotations_(monitoring_annotations) {}
+  KeysetHandle(std::unique_ptr<google::crypto::tink::Keyset> keyset,
+               const std::vector<std::shared_ptr<const Entry>>& entries,
+               const absl::flat_hash_map<std::string, std::string>&
+                   monitoring_annotations)
+      : keyset_(std::move(*keyset)),
+        entries_(entries),
+        monitoring_annotations_(monitoring_annotations) {}
 
   // Generates a key from `key_template` and adds it `keyset`.
   static crypto::tink::util::StatusOr<uint32_t> AddToKeyset(
       const google::crypto::tink::KeyTemplate& key_template, bool as_primary,
       google::crypto::tink::Keyset* keyset);
 
+  // Creates list of KeysetHandle::Entry entries derived from `keyset` in order.
+  static crypto::tink::util::StatusOr<std::vector<std::shared_ptr<const Entry>>>
+  GetEntriesFromKeyset(const google::crypto::tink::Keyset& keyset);
+
+  // Creates KeysetHandle::Entry for `key`, which will be set to primary if
+  // its key id equals `primary_key_id`.
+  static util::StatusOr<Entry> CreateEntry(
+      const google::crypto::tink::Keyset::Key& key, uint32_t primary_key_id);
+
   // Generates a key from `key_template` and adds it to the keyset handle.
   crypto::tink::util::StatusOr<uint32_t> AddKey(
       const google::crypto::tink::KeyTemplate& key_template, bool as_primary);
 
   // Returns keyset held by this handle.
-  const google::crypto::tink::Keyset& get_keyset() const;
+  const google::crypto::tink::Keyset& get_keyset() const { return keyset_; }
 
   // Creates a set of primitives corresponding to the keys with
   // (status == ENABLED) in the keyset given in 'keyset_handle',
@@ -238,7 +276,17 @@ class KeysetHandle {
   crypto::tink::util::StatusOr<std::unique_ptr<PrimitiveSet<P>>> GetPrimitives(
       const KeyManager<P>* custom_manager) const;
 
+  // Creates KeysetHandle::Entry from `keyset_` at `index`.
+  Entry CreateEntryAt(int index) const;
+
   google::crypto::tink::Keyset keyset_;
+  // If this keyset handle has been created with a constructor that does not
+  // accept an entries argument, then `entries` will be empty and operator[]
+  // will fall back to creating the key entry on demand from `keyset_`.
+  //
+  // If `entries_` is not empty, then it should contain exactly one key entry
+  // for each key proto in `keyset_`.
+  std::vector<std::shared_ptr<const Entry>> entries_;
   absl::flat_hash_map<std::string, std::string> monitoring_annotations_;
 };
 
diff --git a/cc/keyset_handle_builder.h b/cc/keyset_handle_builder.h
index 5111ff3f4..72d3af8d8 100644
--- a/cc/keyset_handle_builder.h
+++ b/cc/keyset_handle_builder.h
@@ -57,7 +57,7 @@ class KeysetHandleBuilder {
 
     // Creates new KeysetHandleBuilder::Entry from a given `key`. Also, sets
     // key `status` and whether or not the key `is_primary`.
-    static Entry CreateFromKey(std::unique_ptr<Key> key, KeyStatus status,
+    static Entry CreateFromKey(std::shared_ptr<const Key> key, KeyStatus status,
                                bool is_primary);
 
     template <typename CopyableKey>
@@ -70,7 +70,7 @@ class KeysetHandleBuilder {
     // Creates new KeysetHandleBuilder::Entry from given `parameters`. Also,
     // sets key `status` and whether or not the key `is_primary`. If `id`
     // does not have a value, then the key will be assigned a random id.
-    static Entry CreateFromParams(std::unique_ptr<Parameters> parameters,
+    static Entry CreateFromParams(std::shared_ptr<const Parameters> parameters,
                                   KeyStatus status, bool is_primary,
                                   absl::optional<int> id = absl::nullopt);
 
diff --git a/cc/proto_keyset_format_test.cc b/cc/proto_keyset_format_test.cc
index ebddc3bce..5ac23e60b 100644
--- a/cc/proto_keyset_format_test.cc
+++ b/cc/proto_keyset_format_test.cc
@@ -94,7 +94,7 @@ TEST_F(SerializeKeysetToProtoKeysetFormatTest, SerializeAndParseSingleKey) {
   ASSERT_THAT(handle->size(), Eq(1));
   ASSERT_THAT(parsed_handle->size(), Eq(1));
 
-  EXPECT_TRUE((*handle)[0].GetKey() == (*parsed_handle)[0].GetKey());
+  EXPECT_TRUE(*(*handle)[0].GetKey() == *(*parsed_handle)[0].GetKey());
   EXPECT_TRUE((*handle)[0].GetId() == (*parsed_handle)[0].GetId());
   EXPECT_TRUE((*handle)[0].GetStatus() == (*parsed_handle)[0].GetStatus());
 }
@@ -129,15 +129,15 @@ TEST_F(SerializeKeysetToProtoKeysetFormatTest, SerializeAndParseMultipleKeys) {
   ASSERT_THAT(handle->size(), Eq(3));
   ASSERT_THAT(parsed_handle->size(), Eq(3));
 
-  EXPECT_TRUE((*handle)[0].GetKey() == (*parsed_handle)[0].GetKey());
+  EXPECT_TRUE(*(*handle)[0].GetKey() == *(*parsed_handle)[0].GetKey());
   EXPECT_TRUE((*handle)[0].GetId() == (*parsed_handle)[0].GetId());
   EXPECT_TRUE((*handle)[0].GetStatus() == (*parsed_handle)[0].GetStatus());
 
-  EXPECT_TRUE((*handle)[1].GetKey() == (*parsed_handle)[1].GetKey());
+  EXPECT_TRUE(*(*handle)[1].GetKey() == *(*parsed_handle)[1].GetKey());
   EXPECT_TRUE((*handle)[1].GetId() == (*parsed_handle)[1].GetId());
   EXPECT_TRUE((*handle)[1].GetStatus() == (*parsed_handle)[1].GetStatus());
 
-  EXPECT_TRUE((*handle)[2].GetKey() == (*parsed_handle)[2].GetKey());
+  EXPECT_TRUE(*(*handle)[2].GetKey() == *(*parsed_handle)[2].GetKey());
   EXPECT_TRUE((*handle)[2].GetId() == (*parsed_handle)[2].GetId());
   EXPECT_TRUE((*handle)[2].GetStatus() == (*parsed_handle)[2].GetStatus());
 }
@@ -251,10 +251,11 @@ TEST_F(SerializeKeysetToProtoKeysetFormatTest, SerializeAndParsePublicKey) {
   ASSERT_THAT(parsed_handle3->size(), Eq(1));
   ASSERT_THAT(parsed_handle4->size(), Eq(1));
 
-  EXPECT_TRUE((**public_handle)[0].GetKey() == (*parsed_handle1)[0].GetKey());
-  EXPECT_TRUE((**public_handle)[0].GetKey() == (*parsed_handle2)[0].GetKey());
-  EXPECT_TRUE((**public_handle)[0].GetKey() == (*parsed_handle3)[0].GetKey());
-  EXPECT_TRUE((**public_handle)[0].GetKey() == (*parsed_handle4)[0].GetKey());
+  // TODO(b/277791403): Replace with KeysetHandle::Entry equality checks.
+  EXPECT_TRUE(*(**public_handle)[0].GetKey() == *(*parsed_handle1)[0].GetKey());
+  EXPECT_TRUE(*(**public_handle)[0].GetKey() == *(*parsed_handle2)[0].GetKey());
+  EXPECT_TRUE(*(**public_handle)[0].GetKey() == *(*parsed_handle3)[0].GetKey());
+  EXPECT_TRUE(*(**public_handle)[0].GetKey() == *(*parsed_handle4)[0].GetKey());
 
   EXPECT_TRUE((**public_handle)[0].GetId() == (*parsed_handle1)[0].GetId());
   EXPECT_TRUE((**public_handle)[0].GetId() == (*parsed_handle2)[0].GetId());
diff --git a/objc/Tests/UnitTests/core/TINKCleartextKeysetHandleTest.mm b/objc/Tests/UnitTests/core/TINKCleartextKeysetHandleTest.mm
index f3649f276..12370a11d 100644
--- a/objc/Tests/UnitTests/core/TINKCleartextKeysetHandleTest.mm
+++ b/objc/Tests/UnitTests/core/TINKCleartextKeysetHandleTest.mm
@@ -51,10 +51,10 @@ using ::crypto::tink::util::StatusOr;
 - (void)testReadValidKeyset {
   google::crypto::tink::Keyset keyset;
   google::crypto::tink::Keyset::Key key;
-  crypto::tink::test::AddTinkKey("some key type", 42, key,
+  crypto::tink::test::AddTinkKey("some_key_type", 42, key,
                                  google::crypto::tink::KeyStatusType::ENABLED,
                                  google::crypto::tink::KeyData::SYMMETRIC, &keyset);
-  crypto::tink::test::AddRawKey("some other key type", 711, key,
+  crypto::tink::test::AddRawKey("some_other_key_type", 711, key,
                                 google::crypto::tink::KeyStatusType::ENABLED,
                                 google::crypto::tink::KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
@@ -106,10 +106,10 @@ using ::crypto::tink::util::StatusOr;
 - (void)testSerializeKeyset {
   google::crypto::tink::Keyset keyset;
   google::crypto::tink::Keyset::Key key;
-  crypto::tink::test::AddTinkKey("some key type", 42, key,
+  crypto::tink::test::AddTinkKey("some_key_type", 42, key,
                                  google::crypto::tink::KeyStatusType::ENABLED,
                                  google::crypto::tink::KeyData::SYMMETRIC, &keyset);
-  crypto::tink::test::AddRawKey("some other key type", 711, key,
+  crypto::tink::test::AddRawKey("some_other_key_type", 711, key,
                                 google::crypto::tink::KeyStatusType::ENABLED,
                                 google::crypto::tink::KeyData::SYMMETRIC, &keyset);
   keyset.set_primary_key_id(42);
diff --git a/objc/Tests/UnitTests/core/TINKKeysetHandleTest.mm b/objc/Tests/UnitTests/core/TINKKeysetHandleTest.mm
index c11911383..80d5a6c24 100644
--- a/objc/Tests/UnitTests/core/TINKKeysetHandleTest.mm
+++ b/objc/Tests/UnitTests/core/TINKKeysetHandleTest.mm
@@ -80,10 +80,10 @@ static Keyset *gKeyset;
   gKeyset = new Keyset();
   google::crypto::tink::Keyset::Key ccKey;
 
-  crypto::tink::test::AddTinkKey("some key type", 42, ccKey,
+  crypto::tink::test::AddTinkKey("some_key_type", 42, ccKey,
                                  google::crypto::tink::KeyStatusType::ENABLED,
                                  google::crypto::tink::KeyData::SYMMETRIC, gKeyset);
-  crypto::tink::test::AddRawKey("some other key type", 711, ccKey,
+  crypto::tink::test::AddRawKey("some_other_key_type", 711, ccKey,
                                 google::crypto::tink::KeyStatusType::ENABLED,
                                 google::crypto::tink::KeyData::SYMMETRIC, gKeyset);
   gKeyset->set_primary_key_id(42);
@@ -458,9 +458,9 @@ static Keyset *gKeyset;
 - (void)testReadNoSecret {
   auto keyset = std::make_unique<Keyset>();
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PUBLIC,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PUBLIC,
              keyset.get());
-  AddRawKey("some other key type", 711, key, KeyStatusType::ENABLED, KeyData::REMOTE, keyset.get());
+  AddRawKey("some_other_key_type", 711, key, KeyStatusType::ENABLED, KeyData::REMOTE, keyset.get());
   keyset->set_primary_key_id(42);
   NSData *serializedKeyset = TINKStringToNSData(keyset->SerializeAsString());
   NSError *error = nil;
@@ -480,7 +480,7 @@ static Keyset *gKeyset;
 - (void)testReadNoSecretFailForTypeUnknown {
   auto keyset = std::make_unique<Keyset>();
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED, KeyData::UNKNOWN_KEYMATERIAL,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED, KeyData::UNKNOWN_KEYMATERIAL,
              keyset.get());
   keyset->set_primary_key_id(42);
   NSData *serializedKeyset = TINKStringToNSData(keyset->SerializeAsString());
@@ -497,7 +497,7 @@ static Keyset *gKeyset;
 - (void)testReadNoSecretFailForTypeSymmetric {
   auto keyset = std::make_unique<Keyset>();
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED, KeyData::SYMMETRIC, keyset.get());
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED, KeyData::SYMMETRIC, keyset.get());
   keyset->set_primary_key_id(42);
   NSData *serializedKeyset = TINKStringToNSData(keyset->SerializeAsString());
   NSError *error = nil;
@@ -513,7 +513,7 @@ static Keyset *gKeyset;
 - (void)testReadNoSecretFailForTypeAssymmetricPrivate {
   auto keyset = std::make_unique<Keyset>();
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE,
              keyset.get());
   keyset->set_primary_key_id(42);
   NSData *serializedKeyset = TINKStringToNSData(keyset->SerializeAsString());
@@ -530,13 +530,13 @@ static Keyset *gKeyset;
 - (void)testReadNoSecretFailForHidden {
   auto keyset = std::make_unique<Keyset>();
   Keyset::Key key;
-  AddTinkKey("some key type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PUBLIC,
+  AddTinkKey("some_key_type", 42, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PUBLIC,
              keyset.get());
   for (int i = 0; i < 10; ++i) {
     AddTinkKey(absl::StrCat("more key type", i), i, key, KeyStatusType::ENABLED,
                KeyData::ASYMMETRIC_PUBLIC, keyset.get());
   }
-  AddRawKey("some other key type", 10, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE,
+  AddRawKey("some_other_key_type", 10, key, KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE,
             keyset.get());
   for (int i = 0; i < 10; ++i) {
     AddRawKey(absl::StrCat("more key type", i + 100), i + 100, key, KeyStatusType::ENABLED,