Change internal API RegistryImpl::WrapKeyset to accept annotations.

Notes:
 * The API is internal and only used by KeysetHandle.
 * An empty annotations map indicates the absence of annotations for the Keyset.
PiperOrigin-RevId: 438300092

3 files changed, 18 insertions, 13 deletions

diff --git a/cc/internal/registry_impl.h b/cc/internal/registry_impl.h
index 28548d537..9ac112045 100644
--- a/cc/internal/registry_impl.h
+++ b/cc/internal/registry_impl.h
@@ -135,9 +135,11 @@ class RegistryImpl {
       std::unique_ptr<PrimitiveSet<P>> primitive_set) const
       ABSL_LOCKS_EXCLUDED(maps_mutex_);
 
+  // Wraps a `keyset` and annotates it with `annotations`.
   template <class P>
   crypto::tink::util::StatusOr<std::unique_ptr<P>> WrapKeyset(
-      const google::crypto::tink::Keyset& keyset) const
+      const google::crypto::tink::Keyset& keyset,
+      const absl::flat_hash_map<std::string, std::string>& annotations) const
       ABSL_LOCKS_EXCLUDED(maps_mutex_);
 
   crypto::tink::util::StatusOr<google::crypto::tink::KeyData> DeriveKey(
@@ -792,17 +794,14 @@ crypto::tink::util::StatusOr<std::unique_ptr<P>> RegistryImpl::Wrap(
 
 template <class P>
 crypto::tink::util::StatusOr<std::unique_ptr<P>> RegistryImpl::WrapKeyset(
-    const google::crypto::tink::Keyset& keyset) const {
-  util::StatusOr<const KeysetWrapper<P>*> wrapper_result =
+    const google::crypto::tink::Keyset& keyset,
+    const absl::flat_hash_map<std::string, std::string>& annotations) const {
+  crypto::tink::util::StatusOr<const KeysetWrapper<P>*> keyset_wrapper =
       GetKeysetWrapper<P>();
-  if (!wrapper_result.ok()) {
-    return wrapper_result.status();
+  if (!keyset_wrapper.ok()) {
+    return keyset_wrapper.status();
   }
-  // TODO(b/222245356): Replace empty annotations map with actual annotations
-  // when support is provided to this class.
-  crypto::tink::util::StatusOr<std::unique_ptr<P>> primitive_result =
-      wrapper_result.value()->Wrap(keyset, /*annotations=*/{});
-  return std::move(primitive_result);
+  return (*keyset_wrapper)->Wrap(keyset, annotations);
 }
 
 inline crypto::tink::util::Status RegistryImpl::RestrictToFipsIfEmpty() const {
diff --git a/cc/internal/registry_impl_test.cc b/cc/internal/registry_impl_test.cc
index a5707ce58..f9f459888 100644
--- a/cc/internal/registry_impl_test.cc
+++ b/cc/internal/registry_impl_test.cc
@@ -913,7 +913,8 @@ TEST_F(RegistryTest, KeysetWrappingTest) {
               IsOk());
 
   crypto::tink::util::StatusOr<std::unique_ptr<AeadVariant>> aead_variant =
-      RegistryImpl::GlobalInstance().WrapKeyset<AeadVariant>(keyset);
+      RegistryImpl::GlobalInstance().WrapKeyset<AeadVariant>(
+          keyset, /*annotations=*/{});
   EXPECT_THAT(aead_variant.status(), IsOk());
   EXPECT_THAT(aead_variant.value()->get(), Eq(raw_key));
 }
@@ -937,7 +938,8 @@ TEST_F(RegistryTest, TransformingKeysetWrappingTest) {
               IsOk());
 
   crypto::tink::util::StatusOr<std::unique_ptr<std::string>> string_primitive =
-      RegistryImpl::GlobalInstance().WrapKeyset<std::string>(keyset);
+      RegistryImpl::GlobalInstance().WrapKeyset<std::string>(
+          keyset, /*annotations=*/{});
   EXPECT_THAT(string_primitive.status(), IsOk());
   EXPECT_THAT(*string_primitive.value(), Eq(raw_key));
 }
diff --git a/cc/keyset_handle.h b/cc/keyset_handle.h
index 60296fe6a..817afe11b 100644
--- a/cc/keyset_handle.h
+++ b/cc/keyset_handle.h
@@ -181,7 +181,11 @@ KeysetHandle::GetPrimitives(const KeyManager<P>* custom_manager) const {
 template <class P>
 crypto::tink::util::StatusOr<std::unique_ptr<P>> KeysetHandle::GetPrimitive()
     const {
-  return internal::RegistryImpl::GlobalInstance().WrapKeyset<P>(keyset_);
+  // TODO(b/222245356): Replace second argument with annotations when available
+  // to KeysetHandle via its public interface.
+  return internal::RegistryImpl::GlobalInstance().WrapKeyset<P>(
+      keyset_,
+      /*annotations=*/{});
 }
 
 template <class P>