Implement KeysetHandle::GenerateNew with configs. #tinkApiChange

PiperOrigin-RevId: 542940308

3 files changed, 144 insertions, 84 deletions

diff --git a/cc/core/keyset_handle.cc b/cc/core/keyset_handle.cc
index 94d67232f..eded47316 100644
--- a/cc/core/keyset_handle.cc
+++ b/cc/core/keyset_handle.cc
@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-///////////////////////////////////////////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
 #include "tink/keyset_handle.h"
 
 #include <cstdint>
@@ -29,11 +30,13 @@
 #include "absl/types/optional.h"
 #include "tink/aead.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/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_status.h"
 #include "tink/keyset_reader.h"
 #include "tink/keyset_writer.h"
@@ -336,19 +339,32 @@ util::Status KeysetHandle::WriteNoSecret(KeysetWriter* writer) const {
 }
 
 util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::GenerateNew(
-    const KeyTemplate& key_template,
+    const KeyTemplate& key_template, const KeyGenConfiguration& config,
     const absl::flat_hash_map<std::string, std::string>&
         monitoring_annotations) {
   auto handle =
       absl::WrapUnique(new KeysetHandle(Keyset(), monitoring_annotations));
   util::StatusOr<uint32_t> const result =
-      handle->AddKey(key_template, /*as_primary=*/true);
+      handle->AddKey(key_template, /*as_primary=*/true, config);
   if (!result.ok()) {
     return result.status();
   }
   return std::move(handle);
 }
 
+util::StatusOr<std::unique_ptr<KeysetHandle>> KeysetHandle::GenerateNew(
+    const KeyTemplate& key_template,
+    const absl::flat_hash_map<std::string, std::string>&
+        monitoring_annotations) {
+  KeyGenConfiguration config;
+  util::Status status =
+      internal::KeyGenConfigurationImpl::SetGlobalRegistryMode(config);
+  if (!status.ok()) {
+    return status;
+  }
+  return GenerateNew(key_template, config, monitoring_annotations);
+}
+
 util::StatusOr<std::unique_ptr<Keyset::Key>> ExtractPublicKey(
     const Keyset::Key& key) {
   if (key.key_data().key_material_type() != KeyData::ASYMMETRIC_PRIVATE) {
@@ -389,21 +405,42 @@ KeysetHandle::GetPublicKeysetHandle() const {
 
 crypto::tink::util::StatusOr<uint32_t> KeysetHandle::AddToKeyset(
     const google::crypto::tink::KeyTemplate& key_template, bool as_primary,
-    Keyset* keyset) {
+    const KeyGenConfiguration& config, Keyset* keyset) {
   if (key_template.output_prefix_type() ==
       google::crypto::tink::OutputPrefixType::UNKNOWN_PREFIX) {
     return util::Status(absl::StatusCode::kInvalidArgument,
                         "key template has unknown prefix");
   }
-  auto key_data_result = Registry::NewKeyData(key_template);
-  if (!key_data_result.ok()) return key_data_result.status();
-  auto key_data = std::move(key_data_result.value());
+
+  // Generate new key data.
+  util::StatusOr<std::unique_ptr<KeyData>> key_data;
+  if (internal::KeyGenConfigurationImpl::GetGlobalRegistryMode(config)) {
+    key_data = Registry::NewKeyData(key_template);
+  } 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_template.type_url());
+    if (!key_type_info.ok()) {
+      return key_type_info.status();
+    }
+    key_data = (*key_type_info)->key_factory().NewKeyData(key_template.value());
+  }
+  if (!key_data.ok()) {
+    return key_data.status();
+  }
+
+  // Add and fill in new key in `keyset`.
   Keyset::Key* key = keyset->add_key();
+  *(key->mutable_key_data()) = *std::move(key_data).value();
+  key->set_status(KeyStatusType::ENABLED);
+  key->set_output_prefix_type(key_template.output_prefix_type());
+
   uint32_t key_id = GenerateUnusedKeyId(*keyset);
-  *(key->mutable_key_data()) = *key_data;
-  key->set_status(google::crypto::tink::KeyStatusType::ENABLED);
   key->set_key_id(key_id);
-  key->set_output_prefix_type(key_template.output_prefix_type());
   if (as_primary) {
     keyset->set_primary_key_id(key_id);
   }
@@ -411,8 +448,10 @@ 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) {
-  util::StatusOr<uint32_t> id = AddToKeyset(key_template, as_primary, &keyset_);
+    const google::crypto::tink::KeyTemplate& key_template, bool as_primary,
+    const KeyGenConfiguration& config) {
+  util::StatusOr<uint32_t> id =
+      AddToKeyset(key_template, as_primary, config, &keyset_);
   if (!id.ok()) {
     return id.status();
   }
diff --git a/cc/core/keyset_handle_test.cc b/cc/core/keyset_handle_test.cc
index c381f219f..b28bf697c 100644
--- a/cc/core/keyset_handle_test.cc
+++ b/cc/core/keyset_handle_test.cc
@@ -34,21 +34,19 @@
 #include "tink/binary_keyset_writer.h"
 #include "tink/cleartext_keyset_handle.h"
 #include "tink/config/fips_140_2.h"
+#include "tink/config/key_gen_fips_140_2.h"
 #include "tink/config/tink_config.h"
 #include "tink/core/key_manager_impl.h"
-#include "tink/internal/configuration_impl.h"
 #include "tink/internal/fips_utils.h"
-#include "tink/internal/legacy_proto_parameters.h"
-#include "tink/internal/proto_parameters_serialization.h"
+#include "tink/internal/key_gen_configuration_impl.h"
 #include "tink/json_keyset_reader.h"
 #include "tink/json_keyset_writer.h"
+#include "tink/key_gen_configuration.h"
 #include "tink/key_status.h"
-#include "tink/mac/mac_key_templates.h"
 #include "tink/primitive_set.h"
 #include "tink/primitive_wrapper.h"
 #include "tink/signature/ecdsa_sign_key_manager.h"
 #include "tink/signature/signature_key_templates.h"
-#include "tink/util/protobuf_helper.h"
 #include "tink/util/status.h"
 #include "tink/util/test_keyset_handle.h"
 #include "tink/util/test_matchers.h"
@@ -575,55 +573,94 @@ TEST_F(KeysetHandleTest, WriteEncryptedKeysetWithAssociatedData) {
   EXPECT_EQ(absl::StatusCode::kInvalidArgument, status.code());
 }
 
-TEST_F(KeysetHandleTest, GenerateNewKeysetHandle) {
-  const google::crypto::tink::KeyTemplate* key_templates[] = {
+TEST_F(KeysetHandleTest, GenerateNew) {
+  const google::crypto::tink::KeyTemplate* templates[] = {
       &AeadKeyTemplates::Aes128Gcm(),
       &AeadKeyTemplates::Aes256Gcm(),
       &AeadKeyTemplates::Aes128CtrHmacSha256(),
       &AeadKeyTemplates::Aes256CtrHmacSha256(),
   };
-  for (auto templ : key_templates) {
-    auto handle_result = KeysetHandle::GenerateNew(*templ);
-    EXPECT_TRUE(handle_result.ok())
-        << "Failed for template:\n " << templ->SerializeAsString()
-        << "\n with status: " << handle_result.status();
+  KeyGenConfiguration config;
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::SetGlobalRegistryMode(config),
+              IsOk());
+  for (auto templ : templates) {
+    EXPECT_THAT(KeysetHandle::GenerateNew(*templ).status(), IsOk());
+    EXPECT_THAT(KeysetHandle::GenerateNew(*templ, config).status(), IsOk());
   }
 }
 
+TEST_F(KeysetHandleTest, GenerateNewWithBespokeConfig) {
+  KeyGenConfiguration config;
+  EXPECT_THAT(
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), config).status(),
+      StatusIs(absl::StatusCode::kNotFound));
+
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::AddKeyTypeManager(
+                  absl::make_unique<AesGcmKeyManager>(), config),
+              IsOk());
+  EXPECT_THAT(
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), config).status(),
+      IsOk());
+}
+
+TEST_F(KeysetHandleTest, GenerateNewWithGlobalRegistryConfig) {
+  KeyGenConfiguration config;
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::SetGlobalRegistryMode(config),
+              IsOk());
+  EXPECT_THAT(KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), config),
+              IsOk());
+}
+
 TEST_F(KeysetHandleTest, GenerateNewWithAnnotations) {
   const absl::flat_hash_map<std::string, std::string> kAnnotations = {
       {"key1", "value1"}, {"key2", "value2"}};
 
-  // The template used doesn't make any different w.r.t. annotations.
-  util::StatusOr<std::unique_ptr<KeysetHandle>> keyset_handle =
+  // `handle` depends on the global registry.
+  util::StatusOr<std::unique_ptr<KeysetHandle>> handle =
       KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), kAnnotations);
-  ASSERT_THAT(keyset_handle, IsOk());
-  auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
-  absl::flat_hash_map<std::string, std::string> generated_annotations;
-  EXPECT_CALL(*primitive_wrapper, Wrap(_))
-      .WillOnce(
-          [&generated_annotations](
-              std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
-            generated_annotations = generated_primitive_set->get_annotations();
-            std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
-            return aead;
-          });
-  Registry::Reset();
-  ASSERT_THAT(Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
-              IsOk());
-  ASSERT_THAT(Registry::RegisterKeyTypeManager(
-                  absl::make_unique<FakeAeadKeyManager>(
-                      "type.googleapis.com/google.crypto.tink.AesGcmKey"),
-                  /*new_key_allowed=*/true),
+  ASSERT_THAT(handle, IsOk());
+
+  // `config_handle` uses a config that depends on the global registry.
+  KeyGenConfiguration config;
+  ASSERT_THAT(internal::KeyGenConfigurationImpl::SetGlobalRegistryMode(config),
               IsOk());
+  util::StatusOr<std::unique_ptr<KeysetHandle>> config_handle =
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(), config,
+                                kAnnotations);
+  ASSERT_THAT(config_handle, IsOk());
+
+  for (KeysetHandle h : {**handle, **config_handle}) {
+    auto primitive_wrapper = absl::make_unique<MockAeadPrimitiveWrapper>();
+    absl::flat_hash_map<std::string, std::string> generated_annotations;
+    EXPECT_CALL(*primitive_wrapper, Wrap(_))
+        .WillOnce(
+            [&generated_annotations](
+                std::unique_ptr<PrimitiveSet<Aead>> generated_primitive_set) {
+              generated_annotations =
+                  generated_primitive_set->get_annotations();
+              std::unique_ptr<Aead> aead = absl::make_unique<DummyAead>("");
+              return aead;
+            });
 
-  EXPECT_THAT((*keyset_handle)->GetPrimitive<Aead>(), IsOk());
-  EXPECT_EQ(generated_annotations, kAnnotations);
-  // This is needed to cleanup mocks.
-  Registry::Reset();
+    Registry::Reset();
+    ASSERT_THAT(
+        Registry::RegisterPrimitiveWrapper(std::move(primitive_wrapper)),
+        IsOk());
+    ASSERT_THAT(Registry::RegisterKeyTypeManager(
+                    absl::make_unique<FakeAeadKeyManager>(
+                        "type.googleapis.com/google.crypto.tink.AesGcmKey"),
+                    /*new_key_allowed=*/true),
+                IsOk());
+
+    EXPECT_THAT(h.GetPrimitive<Aead>(), IsOk());
+    EXPECT_EQ(generated_annotations, kAnnotations);
+
+    // This is needed to cleanup mocks.
+    Registry::Reset();
+  }
 }
 
-TEST_F(KeysetHandleTest, GenerateNewKeysetHandleErrors) {
+TEST_F(KeysetHandleTest, GenerateNewErrors) {
   KeyTemplate templ;
   templ.set_type_url("type.googleapis.com/some.unknown.KeyType");
   templ.set_output_prefix_type(OutputPrefixType::TINK);
@@ -783,29 +820,11 @@ TEST_F(KeysetHandleTest, GetPrimitiveWithConfigFips1402) {
     GTEST_SKIP() << "Only test in FIPS mode";
   }
 
-  // TODO(b/265705174): Replace with KeysetHandle::GenerateNew once that takes a
-  // config parameter.
-  KeyTemplate templ = AeadKeyTemplates::Aes128Gcm();
-  util::StatusOr<const internal::KeyTypeInfoStore*> store =
-      internal::ConfigurationImpl::GetKeyTypeInfoStore(ConfigFips140_2());
-  ASSERT_THAT(store, IsOk());
-  util::StatusOr<internal::KeyTypeInfoStore::Info*> info =
-      (*store)->Get(templ.type_url());
-  ASSERT_THAT(info, IsOk());
-
-  util::StatusOr<std::unique_ptr<KeyData>> key_data =
-      (*info)->key_factory().NewKeyData(templ.value());
-  ASSERT_THAT(key_data, IsOk());
-
-  Keyset keyset;
-  uint32_t key_id = 0;
-  test::AddKeyData(**key_data, key_id, OutputPrefixType::TINK,
-                   KeyStatusType::ENABLED, &keyset);
-  keyset.set_primary_key_id(key_id);
-  std::unique_ptr<KeysetHandle> handle =
-      TestKeysetHandle::GetKeysetHandle(keyset);
-
-  EXPECT_THAT(handle->GetPrimitive<Aead>(ConfigFips140_2()), IsOk());
+  util::StatusOr<std::unique_ptr<KeysetHandle>> handle =
+      KeysetHandle::GenerateNew(AeadKeyTemplates::Aes128Gcm(),
+                                KeyGenConfigFips140_2());
+  ASSERT_THAT(handle, IsOk());
+  EXPECT_THAT((*handle)->GetPrimitive<Aead>(ConfigFips140_2()), IsOk());
 }
 
 TEST_F(KeysetHandleTest, GetPrimitiveWithConfigFips1402FailsWithNonFipsHandle) {
diff --git a/cc/core/keyset_manager.cc b/cc/core/keyset_manager.cc
index 6ca880a2e..947b3bd46 100644
--- a/cc/core/keyset_manager.cc
+++ b/cc/core/keyset_manager.cc
@@ -17,14 +17,13 @@
 #include "tink/keyset_manager.h"
 
 #include <memory>
-#include <random>
 #include <utility>
 
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "tink/internal/key_gen_configuration_impl.h"
+#include "tink/key_gen_configuration.h"
 #include "tink/keyset_handle.h"
-#include "tink/keyset_reader.h"
-#include "tink/registry.h"
 #include "tink/util/enums.h"
 #include "tink/util/errors.h"
 #include "tink/util/status.h"
@@ -34,12 +33,12 @@
 namespace crypto {
 namespace tink {
 
+using ::crypto::tink::util::Enums;
+using ::crypto::tink::util::Status;
+using ::crypto::tink::util::StatusOr;
 using google::crypto::tink::Keyset;
 using google::crypto::tink::KeyStatusType;
 using google::crypto::tink::KeyTemplate;
-using crypto::tink::util::Enums;
-using crypto::tink::util::Status;
-using crypto::tink::util::StatusOr;
 
 // static
 StatusOr<std::unique_ptr<KeysetManager>> KeysetManager::New(
@@ -71,17 +70,22 @@ StatusOr<uint32_t> KeysetManager::Add(const KeyTemplate& key_template) {
   return Add(key_template, false);
 }
 
-crypto::tink::util::StatusOr<uint32_t> KeysetManager::Add(
+StatusOr<uint32_t> KeysetManager::Add(
     const google::crypto::tink::KeyTemplate& key_template, bool as_primary) {
+  KeyGenConfiguration config;
+  Status status =
+      internal::KeyGenConfigurationImpl::SetGlobalRegistryMode(config);
+  if (!status.ok()) {
+    return status;
+  }
   absl::MutexLock lock(&keyset_mutex_);
-  return KeysetHandle::AddToKeyset(key_template, as_primary, &keyset_);
+  return KeysetHandle::AddToKeyset(key_template, as_primary, config, &keyset_);
 }
 
 StatusOr<uint32_t> KeysetManager::Rotate(const KeyTemplate& key_template) {
   return Add(key_template, true);
 }
 
-
 Status KeysetManager::Enable(uint32_t key_id) {
   absl::MutexLock lock(&keyset_mutex_);
   for (auto& key : *(keyset_.mutable_key())) {
@@ -129,8 +133,7 @@ Status KeysetManager::Delete(uint32_t key_id) {
                      "Cannot delete primary key (key_id %u).", key_id);
   }
   auto key_field = keyset_.mutable_key();
-  for (auto key_iter = key_field->begin();
-       key_iter != key_field->end();
+  for (auto key_iter = key_field->begin(); key_iter != key_field->end();
        key_iter++) {
     auto key = *key_iter;
     if (key.key_id() == key_id) {
@@ -184,7 +187,6 @@ Status KeysetManager::SetPrimary(uint32_t key_id) {
                    "No key with key_id %u found in the keyset.", key_id);
 }
 
-
 int KeysetManager::KeyCount() const {
   absl::MutexLock lock(&keyset_mutex_);
   return keyset_.key_size();