Add config-dependent KeysetHandle::GetPublicKeysetHandle. #tinkApiChange

PiperOrigin-RevId: 553900744

5 files changed, 190 insertions, 40 deletions

diff --git a/cc/BUILD.bazel b/cc/BUILD.bazel
index f7456ede9..30a622ce0 100644
--- a/cc/BUILD.bazel
+++ b/cc/BUILD.bazel
@@ -482,16 +482,20 @@ cc_library(
         ":keyset_writer",
         ":primitive_set",
         ":registry",
+        "//config:global_registry",
         "//internal:configuration_impl",
         "//internal:key_gen_configuration_impl",
         "//internal:key_info",
         "//internal:key_status_util",
+        "//internal:key_type_info_store",
         "//internal:mutable_serialization_registry",
         "//internal:proto_key_serialization",
         "//internal:util",
         "//proto:tink_cc_proto",
         "//util:errors",
         "//util:keyset_util",
+        "//util:status",
+        "//util:statusor",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:check",
@@ -887,11 +891,14 @@ cc_test(
         "//proto:aes_gcm_siv_cc_proto",
         "//proto:tink_cc_proto",
         "//signature:ecdsa_sign_key_manager",
+        "//signature:ecdsa_verify_key_manager",
         "//signature:signature_key_templates",
         "//util:status",
+        "//util:statusor",
         "//util:test_keyset_handle",
         "//util:test_matchers",
         "//util:test_util",
+        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
diff --git a/cc/CMakeLists.txt b/cc/CMakeLists.txt
index bda6a4023..14d1d1d19 100644
--- a/cc/CMakeLists.txt
+++ b/cc/CMakeLists.txt
@@ -451,15 +451,19 @@ tink_cc_library(
     absl::status
     absl::strings
     absl::optional
+    tink::config::global_registry
     tink::internal::configuration_impl
     tink::internal::key_gen_configuration_impl
     tink::internal::key_info
     tink::internal::key_status_util
+    tink::internal::key_type_info_store
     tink::internal::mutable_serialization_registry
     tink::internal::proto_key_serialization
     tink::internal::util
     tink::util::errors
     tink::util::keyset_util
+    tink::util::status
+    tink::util::statusor
     tink::proto::tink_cc_proto
 )
 
@@ -818,6 +822,7 @@ tink_cc_test(
     tink::core::primitive_set
     tink::core::primitive_wrapper
     gmock
+    absl::memory
     absl::status
     absl::strings
     tink::aead::aead_key_templates
@@ -830,8 +835,10 @@ tink_cc_test(
     tink::internal::fips_utils
     tink::internal::key_gen_configuration_impl
     tink::signature::ecdsa_sign_key_manager
+    tink::signature::ecdsa_verify_key_manager
     tink::signature::signature_key_templates
     tink::util::status
+    tink::util::statusor
     tink::util::test_keyset_handle
     tink::util::test_matchers
     tink::util::test_util
diff --git a/cc/core/keyset_handle.cc b/cc/core/keyset_handle.cc
index 9a6b70c28..9c82423d3 100644
--- a/cc/core/keyset_handle.cc
+++ b/cc/core/keyset_handle.cc
@@ -29,20 +29,25 @@
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 #include "tink/aead.h"
+#include "tink/config/global_registry.h"
 #include "tink/insecure_secret_key_access.h"
 #include "tink/internal/key_gen_configuration_impl.h"
 #include "tink/internal/key_info.h"
 #include "tink/internal/key_status_util.h"
+#include "tink/internal/key_type_info_store.h"
 #include "tink/internal/mutable_serialization_registry.h"
 #include "tink/internal/proto_key_serialization.h"
 #include "tink/internal/util.h"
 #include "tink/key_gen_configuration.h"
+#include "tink/key_manager.h"
 #include "tink/key_status.h"
 #include "tink/keyset_reader.h"
 #include "tink/keyset_writer.h"
 #include "tink/registry.h"
 #include "tink/util/errors.h"
 #include "tink/util/keyset_util.h"
+#include "tink/util/status.h"
+#include "tink/util/statusor.h"
 #include "proto/tink.pb.h"
 
 using google::crypto::tink::EncryptedKeyset;
@@ -366,25 +371,52 @@ util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::GenerateNew(
 }
 
 util::StatusOr<std::unique_ptr<Keyset::Key>> ExtractPublicKey(
-    const Keyset::Key& key) {
+    const Keyset::Key& key, const KeyGenConfiguration& config) {
   if (key.key_data().key_material_type() != KeyData::ASYMMETRIC_PRIVATE) {
     return util::Status(
         absl::StatusCode::kInvalidArgument,
         "Key material is not of type KeyData::ASYMMETRIC_PRIVATE");
   }
-  auto key_data_result = Registry::GetPublicKeyData(key.key_data().type_url(),
-                                                    key.key_data().value());
-  if (!key_data_result.ok()) return key_data_result.status();
+
+  util::StatusOr<std::unique_ptr<KeyData>> key_data;
+  if (internal::KeyGenConfigurationImpl::IsInGlobalRegistryMode(config)) {
+    key_data = Registry::GetPublicKeyData(key.key_data().type_url(),
+                                          key.key_data().value());
+  } else {
+    util::StatusOr<const internal::KeyTypeInfoStore*> key_type_info_store =
+        internal::KeyGenConfigurationImpl::GetKeyTypeInfoStore(config);
+    if (!key_type_info_store.ok()) {
+      return key_type_info_store.status();
+    }
+    util::StatusOr<const internal::KeyTypeInfoStore::Info*> key_type_info =
+        (*key_type_info_store)->Get(key.key_data().type_url());
+    if (!key_type_info.ok()) {
+      return key_type_info.status();
+    }
+    auto factory = dynamic_cast<const PrivateKeyFactory*>(
+        &(*key_type_info)->key_factory());
+    if (factory == nullptr) {
+      return ToStatusF(
+          absl::StatusCode::kInvalidArgument,
+          "KeyManager for type '%s' does not have a PrivateKeyFactory.",
+          key.key_data().type_url());
+    }
+    key_data = factory->GetPublicKeyData(key.key_data().value());
+  }
+  if (!key_data.ok()) {
+    return key_data.status();
+  }
+
   auto public_key = absl::make_unique<Keyset::Key>(key);
-  public_key->mutable_key_data()->Swap(key_data_result.value().get());
+  public_key->mutable_key_data()->Swap(key_data->get());
   return std::move(public_key);
 }
 
 util::StatusOr<std::unique_ptr<KeysetHandle>>
-KeysetHandle::GetPublicKeysetHandle() const {
+KeysetHandle::GetPublicKeysetHandle(const KeyGenConfiguration& config) const {
   std::unique_ptr<Keyset> public_keyset(new Keyset());
   for (const Keyset::Key& key : get_keyset().key()) {
-    auto public_key_result = ExtractPublicKey(key);
+    auto public_key_result = ExtractPublicKey(key, config);
     if (!public_key_result.ok()) return public_key_result.status();
     public_keyset->add_key()->Swap(public_key_result.value().get());
   }
@@ -403,6 +435,11 @@ KeysetHandle::GetPublicKeysetHandle() const {
   return std::move(handle);
 }
 
+util::StatusOr<std::unique_ptr<KeysetHandle>>
+KeysetHandle::GetPublicKeysetHandle() const {
+  return GetPublicKeysetHandle(KeyGenConfigGlobalRegistry());
+}
+
 crypto::tink::util::StatusOr<uint32_t> KeysetHandle::AddToKeyset(
     const google::crypto::tink::KeyTemplate& key_template, bool as_primary,
     const KeyGenConfiguration& config, Keyset* keyset) {
diff --git a/cc/core/keyset_handle_test.cc b/cc/core/keyset_handle_test.cc
index 91aad8eff..1305f7194 100644
--- a/cc/core/keyset_handle_test.cc
+++ b/cc/core/keyset_handle_test.cc
@@ -24,6 +24,7 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "tink/aead/aead_key_templates.h"
@@ -47,8 +48,10 @@
 #include "tink/primitive_set.h"
 #include "tink/primitive_wrapper.h"
 #include "tink/signature/ecdsa_sign_key_manager.h"
+#include "tink/signature/ecdsa_verify_key_manager.h"
 #include "tink/signature/signature_key_templates.h"
 #include "tink/util/status.h"
+#include "tink/util/statusor.h"
 #include "tink/util/test_keyset_handle.h"
 #include "tink/util/test_matchers.h"
 #include "tink/util/test_util.h"
@@ -679,6 +682,42 @@ void CompareKeyMetadata(const Keyset::Key& expected,
   EXPECT_EQ(expected.output_prefix_type(), actual.output_prefix_type());
 }
 
+util::StatusOr<const Keyset> CreateEcdsaMultiKeyset() {
+  Keyset keyset;
+  EcdsaSignKeyManager key_manager;
+  EcdsaKeyFormat key_format;
+
+  if (!key_format.ParseFromString(SignatureKeyTemplates::EcdsaP256().value())) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Failed to parse EcdsaP256 key template");
+  }
+  AddTinkKey(EcdsaSignKeyManager().get_key_type(),
+             /* key_id= */ 623628, key_manager.CreateKey(key_format).value(),
+             KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
+
+  if (!key_format.ParseFromString(
+          SignatureKeyTemplates::EcdsaP384Sha384().value())) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Failed to parse EcdsaP384Sha384 key template");
+  }
+  AddLegacyKey(EcdsaSignKeyManager().get_key_type(),
+               /* key_id= */ 36285, key_manager.CreateKey(key_format).value(),
+               KeyStatusType::DISABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
+
+  if (!key_format.ParseFromString(
+          SignatureKeyTemplates::EcdsaP384Sha512().value())) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Failed to parse EcdsaP384Sha512 key template");
+  }
+  AddRawKey(EcdsaSignKeyManager().get_key_type(),
+            /* key_id= */ 42, key_manager.CreateKey(key_format).value(),
+            KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
+  keyset.set_primary_key_id(42);
+
+  return keyset;
+}
+
+// TODO(b/265865177): Modernize existing GetPublicKeysetHandle tests.
 TEST_F(KeysetHandleTest, GetPublicKeysetHandle) {
   {  // A keyset with a single key.
     auto handle_result =
@@ -697,36 +736,19 @@ TEST_F(KeysetHandleTest, GetPublicKeysetHandle) {
               public_keyset.key(0).key_data().key_material_type());
   }
   {  // A keyset with multiple keys.
-    EcdsaSignKeyManager key_manager;
-    Keyset keyset;
-    int key_count = 3;
-
-    EcdsaKeyFormat key_format;
-    ASSERT_TRUE(
-        key_format.ParseFromString(SignatureKeyTemplates::EcdsaP256().value()));
-    AddTinkKey(EcdsaSignKeyManager().get_key_type(),
-               /* key_id= */ 623628, key_manager.CreateKey(key_format).value(),
-               KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
-    ASSERT_TRUE(key_format.ParseFromString(
-        SignatureKeyTemplates::EcdsaP384Sha384().value()));
-    AddLegacyKey(EcdsaSignKeyManager().get_key_type(),
-                 /* key_id= */ 36285, key_manager.CreateKey(key_format).value(),
-                 KeyStatusType::DISABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
-    ASSERT_TRUE(key_format.ParseFromString(
-        SignatureKeyTemplates::EcdsaP384Sha512().value()));
-    AddRawKey(EcdsaSignKeyManager().get_key_type(),
-              /* key_id= */ 42, key_manager.CreateKey(key_format).value(),
-              KeyStatusType::ENABLED, KeyData::ASYMMETRIC_PRIVATE, &keyset);
-    keyset.set_primary_key_id(42);
-    auto handle = TestKeysetHandle::GetKeysetHandle(keyset);
-    auto public_handle_result = handle->GetPublicKeysetHandle();
-    ASSERT_TRUE(public_handle_result.ok()) << public_handle_result.status();
-    auto public_keyset =
-        TestKeysetHandle::GetKeyset(*(public_handle_result.value()));
-    EXPECT_EQ(keyset.primary_key_id(), public_keyset.primary_key_id());
-    EXPECT_EQ(keyset.key_size(), public_keyset.key_size());
-    for (int i = 0; i < key_count; i++) {
-      CompareKeyMetadata(keyset.key(i), public_keyset.key(i));
+    util::StatusOr<const Keyset> keyset = CreateEcdsaMultiKeyset();
+    ASSERT_THAT(keyset, IsOk());
+    std::unique_ptr<KeysetHandle> handle =
+        TestKeysetHandle::GetKeysetHandle(*keyset);
+    util::StatusOr<std::unique_ptr<KeysetHandle>> public_handle =
+        handle->GetPublicKeysetHandle();
+    ASSERT_THAT(public_handle, IsOk());
+
+    const Keyset& public_keyset = TestKeysetHandle::GetKeyset(**public_handle);
+    EXPECT_EQ(keyset->primary_key_id(), public_keyset.primary_key_id());
+    EXPECT_EQ(keyset->key_size(), public_keyset.key_size());
+    for (int i = 0; i < keyset->key_size(); i++) {
+      CompareKeyMetadata(keyset->key(i), public_keyset.key(i));
       EXPECT_EQ(KeyData::ASYMMETRIC_PUBLIC,
                 public_keyset.key(i).key_data().key_material_type());
     }
@@ -771,6 +793,74 @@ TEST_F(KeysetHandleTest, GetPublicKeysetHandleErrors) {
   }
 }
 
+TEST_F(KeysetHandleTest, GetPublicKeysetHandleWithBespokeConfigSucceeds) {
+  util::StatusOr<const Keyset> keyset = CreateEcdsaMultiKeyset();
+  ASSERT_THAT(keyset, IsOk());
+  std::unique_ptr<KeysetHandle> handle =
+      TestKeysetHandle::GetKeysetHandle(*keyset);
+
+  KeyGenConfiguration config;
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::AddAsymmetricKeyManagers(
+                  absl::make_unique<EcdsaSignKeyManager>(),
+                  absl::make_unique<EcdsaVerifyKeyManager>(), config),
+              IsOk());
+  util::StatusOr<std::unique_ptr<KeysetHandle>> public_handle =
+      handle->GetPublicKeysetHandle(config);
+  ASSERT_THAT(public_handle, IsOk());
+
+  const Keyset& public_keyset = TestKeysetHandle::GetKeyset(**public_handle);
+  EXPECT_EQ(keyset->primary_key_id(), public_keyset.primary_key_id());
+  EXPECT_EQ(keyset->key_size(), public_keyset.key_size());
+  for (int i = 0; i < keyset->key_size(); i++) {
+    CompareKeyMetadata(keyset->key(i), public_keyset.key(i));
+    EXPECT_EQ(KeyData::ASYMMETRIC_PUBLIC,
+              public_keyset.key(i).key_data().key_material_type());
+  }
+}
+
+TEST_F(KeysetHandleTest, GetPublicKeysetHandleWithBespokeConfigFails) {
+  KeyGenConfiguration config;
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::AddKeyTypeManager(
+                  absl::make_unique<AesGcmKeyManager>(), config),
+              IsOk());
+  util::StatusOr<std::unique_ptr<KeysetHandle>> handle =
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), config);
+  ASSERT_THAT(handle, IsOk());
+  EXPECT_THAT((*handle)->GetPublicKeysetHandle(config).status(),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(KeysetHandleTest,
+       GetPublicKeysetHandleWithGlobalRegistryConfigSucceeds) {
+  util::StatusOr<const Keyset> keyset = CreateEcdsaMultiKeyset();
+  ASSERT_THAT(keyset, IsOk());
+  std::unique_ptr<KeysetHandle> handle =
+      TestKeysetHandle::GetKeysetHandle(*keyset);
+
+  util::StatusOr<std::unique_ptr<KeysetHandle>> public_handle =
+      handle->GetPublicKeysetHandle(KeyGenConfigGlobalRegistry());
+  ASSERT_THAT(public_handle, IsOk());
+
+  const Keyset& public_keyset = TestKeysetHandle::GetKeyset(**public_handle);
+  EXPECT_EQ(keyset->primary_key_id(), public_keyset.primary_key_id());
+  EXPECT_EQ(keyset->key_size(), public_keyset.key_size());
+  for (int i = 0; i < keyset->key_size(); i++) {
+    CompareKeyMetadata(keyset->key(i), public_keyset.key(i));
+    EXPECT_EQ(KeyData::ASYMMETRIC_PUBLIC,
+              public_keyset.key(i).key_data().key_material_type());
+  }
+}
+
+TEST_F(KeysetHandleTest, GetPublicKeysetHandleWithGlobalRegistryConfigFails) {
+  util::StatusOr<std::unique_ptr<KeysetHandle>> handle =
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(),
+                                KeyGenConfigGlobalRegistry());
+  ASSERT_THAT(handle, IsOk());
+  EXPECT_THAT(
+      (*handle)->GetPublicKeysetHandle(KeyGenConfigGlobalRegistry()).status(),
+      StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
 TEST_F(KeysetHandleTest, GetPrimitive) {
   Keyset keyset;
   KeyData key_data_0 =
diff --git a/cc/keyset_handle.h b/cc/keyset_handle.h
index a420114c2..3d59b5b93 100644
--- a/cc/keyset_handle.h
+++ b/cc/keyset_handle.h
@@ -39,6 +39,7 @@
 #include "tink/keyset_writer.h"
 #include "tink/primitive_set.h"
 #include "tink/registry.h"
+#include "tink/util/statusor.h"
 #include "proto/tink.pb.h"
 
 namespace crypto {
@@ -174,9 +175,17 @@ class KeysetHandle {
   // `CleartextKeysetHandle`.
   crypto::tink::util::Status WriteNoSecret(KeysetWriter* writer) const;
 
-  // Returns a new KeysetHandle that contains public keys corresponding
-  // to the private keys from this handle.
-  // Returns an error if this handle contains keys that are not private keys.
+  // Returns a new KeysetHandle containing public keys corresponding to the
+  // private keys in this handle. Relies on key type managers stored in `config`
+  // to do so. Returns an error if this handle contains keys that are not
+  // private keys.
+  crypto::tink::util::StatusOr<std::unique_ptr<KeysetHandle>>
+  GetPublicKeysetHandle(const KeyGenConfiguration& config) const;
+
+  // Returns a new KeysetHandle containing public keys corresponding to the
+  // private keys in this handle. Relies on key type managers stored in the
+  // global registry to do so. Returns an error if this handle contains keys
+  // that are not private keys.
   crypto::tink::util::StatusOr<std::unique_ptr<KeysetHandle>>
   GetPublicKeysetHandle() const;