Add proto parser and serializer for SLH-DSA public key.

PiperOrigin-RevId: 625367596

2 files changed, 263 insertions, 8 deletions

diff --git a/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization.cc b/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization.cc
index b7ea80337..c9221dd70 100644
--- a/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization.cc
+++ b/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization.cc
@@ -18,11 +18,20 @@
 
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
 #include "tink/experimental/pqcrypto/signature/slh_dsa_parameters.h"
+#include "tink/experimental/pqcrypto/signature/slh_dsa_public_key.h"
+#include "tink/insecure_secret_key_access.h"
+#include "tink/internal/key_parser.h"
+#include "tink/internal/key_serializer.h"
 #include "tink/internal/mutable_serialization_registry.h"
 #include "tink/internal/parameters_parser.h"
 #include "tink/internal/parameters_serializer.h"
+#include "tink/internal/proto_key_serialization.h"
 #include "tink/internal/proto_parameters_serialization.h"
+#include "tink/partial_key_access.h"
+#include "tink/restricted_data.h"
+#include "tink/secret_key_access_token.h"
 #include "tink/util/status.h"
 #include "tink/util/statusor.h"
 #include "proto/experimental/pqcrypto/slh_dsa.pb.h"
@@ -32,6 +41,7 @@ namespace crypto {
 namespace tink {
 namespace {
 
+using ::google::crypto::tink::KeyData;
 using ::google::crypto::tink::OutputPrefixType;
 using ::google::crypto::tink::SlhDsaHashType;
 using ::google::crypto::tink::SlhDsaKeyFormat;
@@ -44,9 +54,16 @@ using SlhDsaProtoParametersParserImpl =
 using SlhDsaProtoParametersSerializerImpl =
     internal::ParametersSerializerImpl<SlhDsaParameters,
                                        internal::ProtoParametersSerialization>;
+using SlhDsaProtoPublicKeyParserImpl =
+    internal::KeyParserImpl<internal::ProtoKeySerialization, SlhDsaPublicKey>;
+using SlhDsaProtoPublicKeySerializerImpl =
+    internal::KeySerializerImpl<SlhDsaPublicKey,
+                                internal::ProtoKeySerialization>;
 
 const absl::string_view kPrivateTypeUrl =
     "type.googleapis.com/google.crypto.tink.SlhDsaPrivateKey";
+const absl::string_view kPublicTypeUrl =
+    "type.googleapis.com/google.crypto.tink.SlhDsaPublicKey";
 
 util::StatusOr<SlhDsaParameters::Variant> ToVariant(
     OutputPrefixType output_prefix_type) {
@@ -201,6 +218,37 @@ util::StatusOr<SlhDsaParameters> ParseParameters(
                       proto_key_format.params());
 }
 
+util::StatusOr<SlhDsaPublicKey> ParsePublicKey(
+    const internal::ProtoKeySerialization& serialization,
+    absl::optional<SecretKeyAccessToken> token) {
+  if (serialization.TypeUrl() != kPublicTypeUrl) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Wrong type URL when parsing SlhDsaPublicKey.");
+  }
+
+  google::crypto::tink::SlhDsaPublicKey proto_key;
+  const RestrictedData& restricted_data = serialization.SerializedKeyProto();
+  if (!proto_key.ParseFromString(
+          restricted_data.GetSecret(InsecureSecretKeyAccess::Get()))) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Failed to parse SlhDsaPublicKey proto");
+  }
+  if (proto_key.version() != 0) {
+    return util::Status(absl::StatusCode::kInvalidArgument,
+                        "Only version 0 keys are accepted.");
+  }
+
+  util::StatusOr<SlhDsaParameters> parameters =
+      ToParameters(serialization.GetOutputPrefixType(), proto_key.params());
+  if (!parameters.ok()) {
+    return parameters.status();
+  }
+
+  return SlhDsaPublicKey::Create(*parameters, proto_key.key_value(),
+                                 serialization.IdRequirement(),
+                                 GetPartialKeyAccess());
+}
+
 util::StatusOr<internal::ProtoParametersSerialization> SerializeParameters(
     const SlhDsaParameters& parameters) {
   util::StatusOr<OutputPrefixType> output_prefix_type =
@@ -222,16 +270,53 @@ util::StatusOr<internal::ProtoParametersSerialization> SerializeParameters(
       proto_key_format.SerializeAsString());
 }
 
-SlhDsaProtoParametersParserImpl* SlhDsaProtoParametersParser() {
-  static auto* parser =
+util::StatusOr<internal::ProtoKeySerialization> SerializePublicKey(
+    const SlhDsaPublicKey& key, absl::optional<SecretKeyAccessToken> token) {
+  util::StatusOr<SlhDsaParams> params = FromParameters(key.GetParameters());
+  if (!params.ok()) {
+    return params.status();
+  }
+
+  google::crypto::tink::SlhDsaPublicKey proto_key;
+  proto_key.set_version(0);
+  *proto_key.mutable_params() = *params;
+  proto_key.set_key_value(key.GetPublicKeyBytes(GetPartialKeyAccess()));
+
+  util::StatusOr<OutputPrefixType> output_prefix_type =
+      ToOutputPrefixType(key.GetParameters().GetVariant());
+  if (!output_prefix_type.ok()) {
+    return output_prefix_type.status();
+  }
+
+  RestrictedData restricted_output = RestrictedData(
+      proto_key.SerializeAsString(), InsecureSecretKeyAccess::Get());
+  return internal::ProtoKeySerialization::Create(
+      kPublicTypeUrl, restricted_output, KeyData::ASYMMETRIC_PUBLIC,
+      *output_prefix_type, key.GetIdRequirement());
+}
+
+SlhDsaProtoParametersParserImpl& SlhDsaProtoParametersParser() {
+  static auto parser =
       new SlhDsaProtoParametersParserImpl(kPrivateTypeUrl, ParseParameters);
-  return parser;
+  return *parser;
 }
 
-SlhDsaProtoParametersSerializerImpl* SlhDsaProtoParametersSerializer() {
-  static auto* serializer = new SlhDsaProtoParametersSerializerImpl(
+SlhDsaProtoParametersSerializerImpl& SlhDsaProtoParametersSerializer() {
+  static auto serializer = new SlhDsaProtoParametersSerializerImpl(
       kPrivateTypeUrl, SerializeParameters);
-  return serializer;
+  return *serializer;
+}
+
+SlhDsaProtoPublicKeyParserImpl& SlhDsaProtoPublicKeyParser() {
+  static auto* parser =
+      new SlhDsaProtoPublicKeyParserImpl(kPublicTypeUrl, ParsePublicKey);
+  return *parser;
+}
+
+SlhDsaProtoPublicKeySerializerImpl& SlhDsaProtoPublicKeySerializer() {
+  static auto* serializer =
+      new SlhDsaProtoPublicKeySerializerImpl(SerializePublicKey);
+  return *serializer;
 }
 
 }  // namespace
@@ -239,13 +324,26 @@ SlhDsaProtoParametersSerializerImpl* SlhDsaProtoParametersSerializer() {
 util::Status RegisterSlhDsaProtoSerialization() {
   util::Status status =
       internal::MutableSerializationRegistry::GlobalInstance()
-          .RegisterParametersParser(SlhDsaProtoParametersParser());
+          .RegisterParametersParser(&SlhDsaProtoParametersParser());
+  if (!status.ok()) {
+    return status;
+  }
+
+  status =
+      internal::MutableSerializationRegistry::GlobalInstance()
+          .RegisterParametersSerializer(&SlhDsaProtoParametersSerializer());
+  if (!status.ok()) {
+    return status;
+  }
+
+  status = internal::MutableSerializationRegistry::GlobalInstance()
+               .RegisterKeyParser(&SlhDsaProtoPublicKeyParser());
   if (!status.ok()) {
     return status;
   }
 
   return internal::MutableSerializationRegistry::GlobalInstance()
-      .RegisterParametersSerializer(SlhDsaProtoParametersSerializer());
+      .RegisterKeySerializer(&SlhDsaProtoPublicKeySerializer());
 }
 
 }  // namespace tink
diff --git a/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization_test.cc b/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization_test.cc
index 151c255c5..8ca54df8f 100644
--- a/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization_test.cc
+++ b/cc/experimental/pqcrypto/signature/slh_dsa_proto_serialization_test.cc
@@ -25,10 +25,17 @@
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 #include "tink/experimental/pqcrypto/signature/slh_dsa_parameters.h"
+#include "tink/experimental/pqcrypto/signature/slh_dsa_public_key.h"
+#include "tink/insecure_secret_key_access.h"
 #include "tink/internal/mutable_serialization_registry.h"
+#include "tink/internal/proto_key_serialization.h"
 #include "tink/internal/proto_parameters_serialization.h"
 #include "tink/internal/serialization.h"
+#include "tink/key.h"
 #include "tink/parameters.h"
+#include "tink/partial_key_access.h"
+#include "tink/restricted_data.h"
+#include "tink/subtle/random.h"
 #include "tink/util/statusor.h"
 #include "tink/util/test_matchers.h"
 #include "proto/experimental/pqcrypto/slh_dsa.pb.h"
@@ -38,8 +45,10 @@ namespace crypto {
 namespace tink {
 namespace {
 
+using ::crypto::tink::subtle::Random;
 using ::crypto::tink::test::IsOk;
 using ::crypto::tink::test::StatusIs;
+using ::google::crypto::tink::KeyData;
 using ::google::crypto::tink::OutputPrefixType;
 using ::google::crypto::tink::SlhDsaHashType;
 using ::google::crypto::tink::SlhDsaKeyFormat;
@@ -54,6 +63,8 @@ using ::testing::Values;
 
 const absl::string_view kPrivateTypeUrl =
     "type.googleapis.com/google.crypto.tink.SlhDsaPrivateKey";
+const absl::string_view kPublicTypeUrl =
+    "type.googleapis.com/google.crypto.tink.SlhDsaPublicKey";
 
 struct TestCase {
   SlhDsaParameters::Variant variant;
@@ -289,6 +300,152 @@ TEST_P(SlhDsaProtoSerializationTest,
   EXPECT_THAT(key_format.params().key_size(), Eq(64));
 }
 
+TEST_P(SlhDsaProtoSerializationTest, ParsePublicKeyWorks) {
+  TestCase test_case = GetParam();
+  ASSERT_THAT(RegisterSlhDsaProtoSerialization(), IsOk());
+
+  SlhDsaParams params;
+  params.set_sig_type(SlhDsaSignatureType::SMALL_SIGNATURE);
+  params.set_hash_type(SlhDsaHashType::SHA2);
+  params.set_key_size(64);
+
+  std::string raw_key_bytes = Random::GetRandomBytes(32);
+  google::crypto::tink::SlhDsaPublicKey key_proto;
+  key_proto.set_version(0);
+  key_proto.set_key_value(raw_key_bytes);
+  *key_proto.mutable_params() = params;
+  RestrictedData serialized_key = RestrictedData(
+      key_proto.SerializeAsString(), InsecureSecretKeyAccess::Get());
+
+  util::StatusOr<internal::ProtoKeySerialization> serialization =
+      internal::ProtoKeySerialization::Create(
+          kPublicTypeUrl, serialized_key, KeyData::ASYMMETRIC_PUBLIC,
+          test_case.output_prefix_type, test_case.id_requirement);
+  ASSERT_THAT(serialization, IsOk());
+
+  util::StatusOr<std::unique_ptr<Key>> key =
+      internal::MutableSerializationRegistry::GlobalInstance().ParseKey(
+          *serialization, /*token=*/absl::nullopt);
+  ASSERT_THAT(key, IsOk());
+  EXPECT_THAT((*key)->GetIdRequirement(), Eq(test_case.id_requirement));
+  EXPECT_THAT((*key)->GetParameters().HasIdRequirement(),
+              test_case.id_requirement.has_value());
+
+  util::StatusOr<SlhDsaParameters> expected_parameters =
+      SlhDsaParameters::Create(
+          SlhDsaParameters::HashType::kSha2, /*private_key_size_in_bytes=*/64,
+          SlhDsaParameters::SignatureType::kSmallSignature, test_case.variant);
+  ASSERT_THAT(expected_parameters, IsOk());
+
+  util::StatusOr<SlhDsaPublicKey> expected_key =
+      SlhDsaPublicKey::Create(*expected_parameters, raw_key_bytes,
+                              test_case.id_requirement, GetPartialKeyAccess());
+  ASSERT_THAT(expected_key, IsOk());
+
+  EXPECT_THAT(**key, Eq(*expected_key));
+}
+
+TEST_F(SlhDsaProtoSerializationTest,
+       ParsePublicKeyWithInvalidSerializationFails) {
+  ASSERT_THAT(RegisterSlhDsaProtoSerialization(), IsOk());
+
+  RestrictedData serialized_key =
+      RestrictedData("invalid_serialization", InsecureSecretKeyAccess::Get());
+
+  util::StatusOr<internal::ProtoKeySerialization> serialization =
+      internal::ProtoKeySerialization::Create(kPublicTypeUrl, serialized_key,
+                                              KeyData::ASYMMETRIC_PUBLIC,
+                                              OutputPrefixType::TINK,
+                                              /*id_requirement=*/0x23456789);
+  ASSERT_THAT(serialization, IsOk());
+
+  util::StatusOr<std::unique_ptr<Key>> key =
+      internal::MutableSerializationRegistry::GlobalInstance().ParseKey(
+          *serialization, InsecureSecretKeyAccess::Get());
+  EXPECT_THAT(key.status(),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Failed to parse SlhDsaPublicKey proto")));
+}
+
+TEST_F(SlhDsaProtoSerializationTest, ParsePublicKeyWithInvalidVersionFails) {
+  ASSERT_THAT(RegisterSlhDsaProtoSerialization(), IsOk());
+
+  SlhDsaParams params;
+  params.set_sig_type(SlhDsaSignatureType::SMALL_SIGNATURE);
+  params.set_hash_type(SlhDsaHashType::SHA2);
+  params.set_key_size(64);
+
+  std::string raw_key_bytes = Random::GetRandomBytes(32);
+  google::crypto::tink::SlhDsaPublicKey key_proto;
+  key_proto.set_version(1);
+  key_proto.set_key_value(raw_key_bytes);
+  *key_proto.mutable_params() = params;
+  RestrictedData serialized_key = RestrictedData(
+      key_proto.SerializeAsString(), InsecureSecretKeyAccess::Get());
+
+  util::StatusOr<internal::ProtoKeySerialization> serialization =
+      internal::ProtoKeySerialization::Create(kPublicTypeUrl, serialized_key,
+                                              KeyData::ASYMMETRIC_PUBLIC,
+                                              OutputPrefixType::TINK,
+                                              /*id_requirement=*/0x23456789);
+  ASSERT_THAT(serialization, IsOk());
+
+  util::StatusOr<std::unique_ptr<Key>> key =
+      internal::MutableSerializationRegistry::GlobalInstance().ParseKey(
+          *serialization, /*token=*/absl::nullopt);
+  EXPECT_THAT(key.status(),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Only version 0 keys are accepted")));
+}
+
+TEST_P(SlhDsaProtoSerializationTest, SerializePublicKeyWorks) {
+  TestCase test_case = GetParam();
+  ASSERT_THAT(RegisterSlhDsaProtoSerialization(), IsOk());
+
+  util::StatusOr<SlhDsaParameters> parameters = SlhDsaParameters::Create(
+      SlhDsaParameters::HashType::kSha2, /*private_key_size_in_bytes=*/64,
+      SlhDsaParameters::SignatureType::kSmallSignature, test_case.variant);
+  ASSERT_THAT(parameters, IsOk());
+
+  std::string raw_key_bytes = Random::GetRandomBytes(32);
+  util::StatusOr<SlhDsaPublicKey> key =
+      SlhDsaPublicKey::Create(*parameters, raw_key_bytes,
+                              test_case.id_requirement, GetPartialKeyAccess());
+  ASSERT_THAT(key, IsOk());
+
+  util::StatusOr<std::unique_ptr<Serialization>> serialization =
+      internal::MutableSerializationRegistry::GlobalInstance()
+          .SerializeKey<internal::ProtoKeySerialization>(
+              *key, /*token=*/absl::nullopt);
+  ASSERT_THAT(serialization, IsOk());
+  EXPECT_THAT((*serialization)->ObjectIdentifier(), Eq(kPublicTypeUrl));
+
+  const internal::ProtoKeySerialization* proto_serialization =
+      dynamic_cast<const internal::ProtoKeySerialization*>(
+          serialization->get());
+  ASSERT_THAT(proto_serialization, NotNull());
+  EXPECT_THAT(proto_serialization->TypeUrl(), Eq(kPublicTypeUrl));
+  EXPECT_THAT(proto_serialization->KeyMaterialType(),
+              Eq(KeyData::ASYMMETRIC_PUBLIC));
+  EXPECT_THAT(proto_serialization->GetOutputPrefixType(),
+              Eq(test_case.output_prefix_type));
+  EXPECT_THAT(proto_serialization->IdRequirement(),
+              Eq(test_case.id_requirement));
+
+  google::crypto::tink::SlhDsaPublicKey proto_key;
+  ASSERT_THAT(proto_key.ParseFromString(
+                  proto_serialization->SerializedKeyProto().GetSecret(
+                      InsecureSecretKeyAccess::Get())),
+              IsTrue());
+  EXPECT_THAT(proto_key.version(), Eq(0));
+  EXPECT_THAT(proto_key.key_value(), Eq(raw_key_bytes));
+  EXPECT_THAT(proto_key.has_params(), IsTrue());
+  EXPECT_THAT(proto_key.params().key_size(), Eq(64));
+  EXPECT_THAT(proto_key.params().hash_type(), Eq(SlhDsaHashType::SHA2));
+  EXPECT_THAT(proto_key.params().sig_type(),
+              Eq(SlhDsaSignatureType::SMALL_SIGNATURE));
+}
+
 }  // namespace
 }  // namespace tink
 }  // namespace crypto