Add get_all_in_keyset_order() to C++ PrimitiveSet. #tinkApiChange

The now deleted KeysetDeriverSetWrapperImpl class allowed KeysetDeriverSetWrapper access to a PrimitiveSet's entries in keyset key order. With the addition of get_all_in_keyset_order() to PrimitiveSet's public API, KeysetDeriverSetWrapperImpl is no longer necessary.

PiperOrigin-RevId: 540170074

8 files changed, 40 insertions, 244 deletions

diff --git a/cc/keyderivation/BUILD.bazel b/cc/keyderivation/BUILD.bazel
index 790a98517..449698a1d 100644
--- a/cc/keyderivation/BUILD.bazel
+++ b/cc/keyderivation/BUILD.bazel
@@ -94,7 +94,6 @@ cc_library(
         "//:cleartext_keyset_handle",
         "//:primitive_set",
         "//:primitive_wrapper",
-        "//keyderivation/internal:keyset_deriver_set_wrapper_impl",
         "//proto:tink_cc_proto",
         "@com_google_absl//absl/status",
     ],
@@ -104,8 +103,10 @@ cc_test(
     name = "keyset_deriver_wrapper_test",
     srcs = ["keyset_deriver_wrapper_test.cc"],
     deps = [
+        ":keyset_deriver",
         ":keyset_deriver_wrapper",
         "//:cleartext_keyset_handle",
+        "//:primitive_set",
         "//proto:tink_cc_proto",
         "//util:test_matchers",
         "@com_google_absl//absl/status",
diff --git a/cc/keyderivation/CMakeLists.txt b/cc/keyderivation/CMakeLists.txt
index 7b363a1ed..fa3d7e849 100644
--- a/cc/keyderivation/CMakeLists.txt
+++ b/cc/keyderivation/CMakeLists.txt
@@ -90,7 +90,6 @@ tink_cc_library(
     tink::core::cleartext_keyset_handle
     tink::core::primitive_set
     tink::core::primitive_wrapper
-    tink::keyderivation::internal::keyset_deriver_set_wrapper_impl
     tink::proto::tink_cc_proto
 )
 
@@ -99,10 +98,12 @@ tink_cc_test(
   SRCS
     keyset_deriver_wrapper_test.cc
   DEPS
+    tink::keyderivation::keyset_deriver
     tink::keyderivation::keyset_deriver_wrapper
     gmock
     absl::status
     tink::core::cleartext_keyset_handle
+    tink::core::primitive_set
     tink::util::test_matchers
     tink::proto::tink_cc_proto
 )
diff --git a/cc/keyderivation/internal/BUILD.bazel b/cc/keyderivation/internal/BUILD.bazel
index f4cce8cb6..09aa90b28 100644
--- a/cc/keyderivation/internal/BUILD.bazel
+++ b/cc/keyderivation/internal/BUILD.bazel
@@ -3,30 +3,6 @@ package(default_visibility = ["//:__subpackages__"])
 licenses(["notice"])
 
 cc_library(
-    name = "keyset_deriver_set_wrapper_impl",
-    hdrs = ["keyset_deriver_set_wrapper_impl.h"],
-    include_prefix = "tink/keyderivation/internal",
-    deps = [
-        "//:primitive_set",
-        "//keyderivation:keyset_deriver",
-    ],
-)
-
-cc_test(
-    name = "keyset_deriver_set_wrapper_impl_test",
-    srcs = ["keyset_deriver_set_wrapper_impl_test.cc"],
-    deps = [
-        ":keyset_deriver_set_wrapper_impl",
-        "//:cleartext_keyset_handle",
-        "//:crypto_format",
-        "//:primitive_set",
-        "//keyderivation:keyset_deriver",
-        "//util:test_matchers",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-cc_library(
     name = "prf_based_deriver",
     srcs = ["prf_based_deriver.cc"],
     hdrs = ["prf_based_deriver.h"],
diff --git a/cc/keyderivation/internal/CMakeLists.txt b/cc/keyderivation/internal/CMakeLists.txt
index 6b01fdecc..2fd4c7ed4 100644
--- a/cc/keyderivation/internal/CMakeLists.txt
+++ b/cc/keyderivation/internal/CMakeLists.txt
@@ -1,29 +1,6 @@
 tink_module(keyderivation::internal)
 
 tink_cc_library(
-  NAME keyset_deriver_set_wrapper_impl
-  SRCS
-    keyset_deriver_set_wrapper_impl.h
-  DEPS
-    tink::core::primitive_set
-    tink::keyderivation::keyset_deriver
-)
-
-tink_cc_test(
-  NAME keyset_deriver_set_wrapper_impl_test
-  SRCS
-    keyset_deriver_set_wrapper_impl_test.cc
-  DEPS
-    tink::keyderivation::internal::keyset_deriver_set_wrapper_impl
-    gmock
-    tink::core::cleartext_keyset_handle
-    tink::core::crypto_format
-    tink::core::primitive_set
-    tink::keyderivation::keyset_deriver
-    tink::util::test_matchers
-)
-
-tink_cc_library(
   NAME prf_based_deriver
   SRCS
     prf_based_deriver.cc
diff --git a/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl.h b/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl.h
deleted file mode 100644
index 63e9304a8..000000000
--- a/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl.h
+++ /dev/null
@@ -1,46 +0,0 @@
-// Copyright 2023 Google Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-//
-////////////////////////////////////////////////////////////////////////////////
-
-#ifndef TINK_KEYDERIVATION_INTERNAL_KEYSET_DERIVER_SET_WRAPPER_IMPL_H_
-#define TINK_KEYDERIVATION_INTERNAL_KEYSET_DERIVER_SET_WRAPPER_IMPL_H_
-
-#include <vector>
-
-#include "tink/keyderivation/keyset_deriver.h"
-#include "tink/primitive_set.h"
-
-namespace crypto {
-namespace tink {
-namespace internal {
-
-class KeysetDeriverSetWrapperImpl {
- public:
-  // Stores PrfBasedDeriverKey entries so key derivation preserves the original
-  // keyset key order.
-  static inline std::vector<PrimitiveSet<KeysetDeriver>::Entry<KeysetDeriver>*>
-  get_all_in_keyset_order(const PrimitiveSet<KeysetDeriver>& primitive_set) {
-    absl::MutexLockMaybe lock(primitive_set.primitives_mutex_.get());
-    std::vector<PrimitiveSet<KeysetDeriver>::Entry<KeysetDeriver>*> result =
-        primitive_set.ordered_keyset_deriver_primitives_;
-    return result;
-  }
-};
-
-}  // namespace internal
-}  // namespace tink
-}  // namespace crypto
-
-#endif  // TINK_KEYDERIVATION_INTERNAL_KEYSET_DERIVER_SET_WRAPPER_IMPL_H_
diff --git a/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl_test.cc b/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl_test.cc
deleted file mode 100644
index 1edad523f..000000000
--- a/cc/keyderivation/internal/keyset_deriver_set_wrapper_impl_test.cc
+++ /dev/null
@@ -1,115 +0,0 @@
-// Copyright 2023 Google Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-//
-////////////////////////////////////////////////////////////////////////////////
-
-#include "tink/keyderivation/internal/keyset_deriver_set_wrapper_impl.h"
-
-#include <memory>
-#include <string>
-#include <vector>
-
-#include "gmock/gmock.h"
-#include "gtest/gtest.h"
-#include "tink/cleartext_keyset_handle.h"
-#include "tink/crypto_format.h"
-#include "tink/keyderivation/keyset_deriver.h"
-#include "tink/primitive_set.h"
-#include "tink/util/test_matchers.h"
-
-namespace crypto {
-namespace tink {
-namespace internal {
-namespace {
-
-using ::crypto::tink::test::IsOk;
-using ::google::crypto::tink::Keyset;
-using ::google::crypto::tink::KeysetInfo;
-using ::google::crypto::tink::KeyStatusType;
-using ::google::crypto::tink::OutputPrefixType;
-using ::testing::Eq;
-using ::testing::SizeIs;
-
-// TODO(b/255828521): Move this to a shared location once KeysetDeriver is in
-// the public API.
-class DummyDeriver : public KeysetDeriver {
- public:
-  explicit DummyDeriver() = default;
-  crypto::tink::util::StatusOr<std::unique_ptr<KeysetHandle>> DeriveKeyset(
-      absl::string_view salt) const override {
-    Keyset keyset;
-    return CleartextKeysetHandle::GetKeysetHandle(keyset);
-  }
-};
-
-TEST(KeysetDeriverSetWrapperImpl, GetAllInKeysetOrder) {
-  auto pset = absl::make_unique<PrimitiveSet<KeysetDeriver>>();
-  std::vector<KeysetInfo::KeyInfo> key_infos;
-
-  KeysetInfo::KeyInfo key_info;
-  key_info.set_key_id(1010101);
-  key_info.set_status(KeyStatusType::ENABLED);
-  key_info.set_output_prefix_type(OutputPrefixType::RAW);
-  key_info.set_type_url(
-      "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  ASSERT_THAT(pset->AddPrimitive(absl::make_unique<DummyDeriver>(), key_info),
-              IsOk());
-  key_infos.push_back(key_info);
-
-  key_info.set_key_id(2020202);
-  key_info.set_status(KeyStatusType::ENABLED);
-  key_info.set_output_prefix_type(OutputPrefixType::LEGACY);
-  key_info.set_type_url(
-      "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  ASSERT_THAT(pset->AddPrimitive(absl::make_unique<DummyDeriver>(), key_info),
-              IsOk());
-  key_infos.push_back(key_info);
-
-  key_info.set_key_id(3030303);
-  key_info.set_status(KeyStatusType::ENABLED);
-  key_info.set_output_prefix_type(OutputPrefixType::TINK);
-  key_info.set_type_url(
-      "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  ASSERT_THAT(pset->AddPrimitive(absl::make_unique<DummyDeriver>(), key_info),
-              IsOk());
-  key_infos.push_back(key_info);
-
-  // Should not be returned by get_all_in_keyset_order() because the type URL is
-  // not PrfBasedDeriverKey.
-  key_info.set_key_id(4040404);
-  key_info.set_status(KeyStatusType::ENABLED);
-  key_info.set_output_prefix_type(OutputPrefixType::TINK);
-  key_info.set_type_url("type.googleapis.com/google.crypto.tink.AesGcmKey");
-  ASSERT_THAT(pset->AddPrimitive(absl::make_unique<DummyDeriver>(), key_info),
-              IsOk());
-
-  std::vector<PrimitiveSet<KeysetDeriver>::Entry<KeysetDeriver>*> entries =
-      KeysetDeriverSetWrapperImpl::get_all_in_keyset_order(*pset);
-  EXPECT_THAT(entries, SizeIs(key_infos.size()));
-
-  for (int i = 0; i < entries.size(); i++) {
-    EXPECT_THAT(entries[i]->get_identifier(),
-                Eq(*CryptoFormat::GetOutputPrefix(key_infos[i])));
-    EXPECT_THAT(entries[i]->get_status(), Eq(KeyStatusType::ENABLED));
-    EXPECT_THAT(entries[i]->get_key_id(), Eq(key_infos[i].key_id()));
-    EXPECT_THAT(entries[i]->get_output_prefix_type(),
-                Eq(key_infos[i].output_prefix_type()));
-    EXPECT_THAT(entries[i]->get_key_type_url(), Eq(key_infos[i].type_url()));
-  }
-}
-
-}  // namespace
-}  // namespace internal
-}  // namespace tink
-}  // namespace crypto
diff --git a/cc/keyderivation/keyset_deriver_wrapper.cc b/cc/keyderivation/keyset_deriver_wrapper.cc
index 6a71a02f2..cf74a1103 100644
--- a/cc/keyderivation/keyset_deriver_wrapper.cc
+++ b/cc/keyderivation/keyset_deriver_wrapper.cc
@@ -21,7 +21,6 @@
 
 #include "absl/status/status.h"
 #include "tink/cleartext_keyset_handle.h"
-#include "tink/keyderivation/internal/keyset_deriver_set_wrapper_impl.h"
 #include "proto/tink.pb.h"
 
 namespace crypto {
@@ -76,9 +75,7 @@ crypto::tink::util::StatusOr<KeyData> DeriveAndGetKeyData(
 crypto::tink::util::StatusOr<std::unique_ptr<KeysetHandle>>
 KeysetDeriverSetWrapper::DeriveKeyset(absl::string_view salt) const {
   Keyset keyset;
-  for (const auto* entry :
-       internal::KeysetDeriverSetWrapperImpl::get_all_in_keyset_order(
-           *deriver_set_)) {
+  for (const auto* entry : deriver_set_->get_all_in_keyset_order()) {
     Keyset::Key* key = keyset.add_key();
 
     crypto::tink::util::StatusOr<KeyData> key_data_or =
diff --git a/cc/keyderivation/keyset_deriver_wrapper_test.cc b/cc/keyderivation/keyset_deriver_wrapper_test.cc
index 6c2aee219..dde29fbcf 100644
--- a/cc/keyderivation/keyset_deriver_wrapper_test.cc
+++ b/cc/keyderivation/keyset_deriver_wrapper_test.cc
@@ -19,11 +19,14 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "absl/status/status.h"
 #include "tink/cleartext_keyset_handle.h"
+#include "tink/keyderivation/keyset_deriver.h"
+#include "tink/primitive_set.h"
 #include "tink/util/test_matchers.h"
 #include "proto/tink.pb.h"
 
@@ -45,7 +48,7 @@ using ::testing::HasSubstr;
 class DummyDeriver : public KeysetDeriver {
  public:
   explicit DummyDeriver(absl::string_view name) : name_(name) {}
-  crypto::tink::util::StatusOr<std::unique_ptr<KeysetHandle>> DeriveKeyset(
+  util::StatusOr<std::unique_ptr<KeysetHandle>> DeriveKeyset(
       absl::string_view salt) const override {
     Keyset::Key key;
     key.mutable_key_data()->set_type_url(
@@ -129,7 +132,8 @@ TEST(KeysetDeriverWrapperTest, WrapSingle) {
 }
 
 TEST(KeysetDeriverWrapperTest, WrapMultiple) {
-  auto deriver_set = absl::make_unique<PrimitiveSet<KeysetDeriver>>();
+  auto pset = absl::make_unique<PrimitiveSet<KeysetDeriver>>();
+  std::vector<KeysetInfo::KeyInfo> key_infos;
 
   KeysetInfo::KeyInfo key_info;
   key_info.set_key_id(1010101);
@@ -137,54 +141,55 @@ TEST(KeysetDeriverWrapperTest, WrapMultiple) {
   key_info.set_output_prefix_type(OutputPrefixType::RAW);
   key_info.set_type_url(
       "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  EXPECT_THAT(
-      deriver_set->AddPrimitive(absl::make_unique<DummyDeriver>("k1"), key_info)
+  ASSERT_THAT(
+      pset->AddPrimitive(absl::make_unique<DummyDeriver>("k1"), key_info)
           .status(),
       IsOk());
+  key_infos.push_back(key_info);
 
   key_info.set_key_id(2020202);
   key_info.set_status(KeyStatusType::ENABLED);
   key_info.set_output_prefix_type(OutputPrefixType::LEGACY);
   key_info.set_type_url(
       "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  auto entry_or = deriver_set->AddPrimitive(
-      absl::make_unique<DummyDeriver>("k2"), key_info);
-  ASSERT_THAT(entry_or, IsOk());
-  EXPECT_THAT(deriver_set->set_primary(entry_or.value()), IsOk());
+  util::StatusOr<PrimitiveSet<KeysetDeriver>::Entry<KeysetDeriver>*> entry =
+      pset->AddPrimitive(absl::make_unique<DummyDeriver>("k2"), key_info);
+  ASSERT_THAT(entry, IsOk());
+  ASSERT_THAT(pset->set_primary(*entry), IsOk());
+  key_infos.push_back(key_info);
 
   key_info.set_key_id(3030303);
   key_info.set_status(KeyStatusType::ENABLED);
   key_info.set_output_prefix_type(OutputPrefixType::TINK);
   key_info.set_type_url(
       "type.googleapis.com/google.crypto.tink.PrfBasedDeriverKey");
-  entry_or = deriver_set->AddPrimitive(absl::make_unique<DummyDeriver>("k3"),
-                                       key_info);
-
-  auto wrapper_deriver_or = KeysetDeriverWrapper().Wrap(std::move(deriver_set));
-  ASSERT_THAT(wrapper_deriver_or, IsOk());
+  ASSERT_THAT(
+      pset->AddPrimitive(absl::make_unique<DummyDeriver>("k3"), key_info),
+      IsOk());
+  key_infos.push_back(key_info);
 
-  auto derived_keyset_or = wrapper_deriver_or.value()->DeriveKeyset("salt");
-  ASSERT_THAT(derived_keyset_or, IsOk());
+  util::StatusOr<std::unique_ptr<KeysetDeriver>> wrapper_deriver =
+      KeysetDeriverWrapper().Wrap(std::move(pset));
+  ASSERT_THAT(wrapper_deriver, IsOk());
 
-  Keyset keyset = CleartextKeysetHandle::GetKeyset(*derived_keyset_or.value());
+  util::StatusOr<std::unique_ptr<KeysetHandle>> derived_keyset =
+      (*wrapper_deriver)->DeriveKeyset("salt");
+  ASSERT_THAT(derived_keyset, IsOk());
+  Keyset keyset = CleartextKeysetHandle::GetKeyset(**derived_keyset);
 
   EXPECT_THAT(keyset.primary_key_id(), Eq(2020202));
   ASSERT_THAT(keyset.key_size(), Eq(3));
 
-  EXPECT_THAT(keyset.key(0).key_data().type_url(), Eq("2:k1salt"));
-  EXPECT_THAT(keyset.key(0).status(), Eq(KeyStatusType::ENABLED));
-  EXPECT_THAT(keyset.key(0).key_id(), Eq(1010101));
-  EXPECT_THAT(keyset.key(0).output_prefix_type(), Eq(OutputPrefixType::RAW));
-
-  EXPECT_THAT(keyset.key(1).key_data().type_url(), Eq("2:k2salt"));
-  EXPECT_THAT(keyset.key(1).status(), Eq(KeyStatusType::ENABLED));
-  EXPECT_THAT(keyset.key(1).key_id(), Eq(2020202));
-  EXPECT_THAT(keyset.key(1).output_prefix_type(), Eq(OutputPrefixType::LEGACY));
-
-  EXPECT_THAT(keyset.key(2).key_data().type_url(), Eq("2:k3salt"));
-  EXPECT_THAT(keyset.key(2).status(), Eq(KeyStatusType::ENABLED));
-  EXPECT_THAT(keyset.key(2).key_id(), Eq(3030303));
-  EXPECT_THAT(keyset.key(2).output_prefix_type(), Eq(OutputPrefixType::TINK));
+  for (int i = 0; i < keyset.key().size(); i++) {
+    std::string type_url = absl::StrCat("2:k", i + 1, "salt");
+    EXPECT_THAT(keyset.key(i).key_data().type_url(), Eq(type_url));
+
+    Keyset::Key key = keyset.key(i);
+    key_info = key_infos[i];
+    EXPECT_THAT(key.status(), Eq(key_info.status()));
+    EXPECT_THAT(key.key_id(), Eq(key_info.key_id()));
+    EXPECT_THAT(key.output_prefix_type(), Eq(key_info.output_prefix_type()));
+  }
 }
 
 }  // namespace