Add DILITHIUM_public_from_private function.

Change-Id: I9b9c53d4f994bd0fc731c9fff3119f0d388900e9
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/69667
Reviewed-by: Adam Langley <agl@google.com>
Commit-Queue: Adam Langley <agl@google.com>

crypto/dilithium/dilithium.c

diff --git a/crypto/dilithium/dilithium.c b/crypto/dilithium/dilithium.c
index 4f209bf..9630715 100644
--- a/crypto/dilithium/dilithium.c
+++ b/crypto/dilithium/dilithium.c
@@ -1214,6 +1214,48 @@
   return ret;
 }
 
+int DILITHIUM_public_from_private(
+    struct DILITHIUM_public_key *out_public_key,
+    const struct DILITHIUM_private_key *private_key) {
+  int ret = 0;
+
+  // Intermediate values, allocated on the heap to allow use when there is a
+  // limited amount of stack.
+  struct values_st {
+    matrix a_ntt;
+    vectorl s1_ntt;
+    vectork t;
+    vectork t0;
+  };
+  struct values_st *values = OPENSSL_malloc(sizeof(*values));
+  if (values == NULL) {
+    goto err;
+  }
+
+  const struct private_key *priv = private_key_from_external(private_key);
+  struct public_key *pub = public_key_from_external(out_public_key);
+
+  OPENSSL_memcpy(pub->rho, priv->rho, sizeof(pub->rho));
+  OPENSSL_memcpy(pub->public_key_hash, priv->public_key_hash,
+                 sizeof(pub->public_key_hash));
+
+  matrix_expand(&values->a_ntt, priv->rho);
+
+  OPENSSL_memcpy(&values->s1_ntt, &priv->s1, sizeof(values->s1_ntt));
+  vectorl_ntt(&values->s1_ntt);
+
+  matrix_mult(&values->t, &values->a_ntt, &values->s1_ntt);
+  vectork_inverse_ntt(&values->t);
+  vectork_add(&values->t, &values->t, &priv->s2);
+
+  vectork_power2_round(&pub->t1, &values->t0, &values->t);
+
+  ret = 1;
+err:
+  OPENSSL_free(values);
+  return ret;
+}
+
 // FIPS 204, Algorithm 2 (`ML-DSA.Sign`). Returns 1 on success and 0 on failure.
 static int dilithium_sign_with_randomizer(
     uint8_t out_encoded_signature[DILITHIUM_SIGNATURE_BYTES],

crypto/dilithium/dilithium_test.cc

diff --git a/crypto/dilithium/dilithium_test.cc b/crypto/dilithium/dilithium_test.cc
index dcd2e8b..51cb295 100644
--- a/crypto/dilithium/dilithium_test.cc
+++ b/crypto/dilithium/dilithium_test.cc
@@ -92,6 +92,23 @@
             1);
 }
 
+TEST(DilithiumTest, PublicFromPrivateIsConsistent) {
+  std::vector<uint8_t> encoded_public_key(DILITHIUM_PUBLIC_KEY_BYTES);
+  auto priv = std::make_unique<DILITHIUM_private_key>();
+  EXPECT_TRUE(DILITHIUM_generate_key(encoded_public_key.data(), priv.get()));
+
+  auto pub = std::make_unique<DILITHIUM_public_key>();
+  EXPECT_TRUE(DILITHIUM_public_from_private(pub.get(), priv.get()));
+
+  std::vector<uint8_t> encoded_public_key2(DILITHIUM_PUBLIC_KEY_BYTES);
+
+  CBB cbb;
+  CBB_init_fixed(&cbb, encoded_public_key2.data(), encoded_public_key2.size());
+  ASSERT_TRUE(DILITHIUM_marshal_public_key(&cbb, pub.get()));
+
+  EXPECT_EQ(Bytes(encoded_public_key2), Bytes(encoded_public_key));
+}
+
 TEST(DilithiumTest, InvalidPublicKeyEncodingLength) {
   // Encode a public key with a trailing 0 at the end.
   std::vector<uint8_t> encoded_public_key(DILITHIUM_PUBLIC_KEY_BYTES + 1);
@@ -100,13 +117,13 @@
 
   // Public key is 1 byte too short.
   CBS cbs = bssl::MakeConstSpan(encoded_public_key)
-            .first(DILITHIUM_PUBLIC_KEY_BYTES - 1);
+                .first(DILITHIUM_PUBLIC_KEY_BYTES - 1);
   auto parsed_pub = std::make_unique<DILITHIUM_public_key>();
   EXPECT_FALSE(DILITHIUM_parse_public_key(parsed_pub.get(), &cbs));
 
   // Public key has the correct length.
-  cbs = bssl::MakeConstSpan(encoded_public_key)
-            .first(DILITHIUM_PUBLIC_KEY_BYTES);
+  cbs =
+      bssl::MakeConstSpan(encoded_public_key).first(DILITHIUM_PUBLIC_KEY_BYTES);
   EXPECT_TRUE(DILITHIUM_parse_public_key(parsed_pub.get(), &cbs));
 
   // Public key is 1 byte too long.
@@ -142,10 +159,6 @@
   EXPECT_FALSE(DILITHIUM_parse_private_key(parsed_priv.get(), &cbs));
 }
 
-// TODO: Add more parsing tests:
-// - signed values out of range (private key's s1/s2 beyond ETA)
-// - signature hints not in the canonical order (signature malleability)
-
 static void DilithiumFileTest(FileTest *t) {
   std::vector<uint8_t> seed, message, public_key_expected, private_key_expected,
       signed_message_expected;

include/openssl/experimental/dilithium.h

diff --git a/include/openssl/experimental/dilithium.h b/include/openssl/experimental/dilithium.h
index 72e79bd..2273421 100644
--- a/include/openssl/experimental/dilithium.h
+++ b/include/openssl/experimental/dilithium.h
@@ -70,6 +70,12 @@
     uint8_t out_encoded_public_key[DILITHIUM_PUBLIC_KEY_BYTES],
     struct DILITHIUM_private_key *out_private_key);
 
+// DILITHIUM_public_from_private sets |*out_public_key| to the public key that
+// corresponds to |private_key|. Returns 1 on success and 0 on failure.
+OPENSSL_EXPORT int DILITHIUM_public_from_private(
+    struct DILITHIUM_public_key *out_public_key,
+    const struct DILITHIUM_private_key *private_key);
+
 // DILITHIUM_sign generates a signature for the message |msg| of length
 // |msg_len| using |private_key| following the randomized algorithm, and writes
 // the encoded signature to |out_encoded_signature|. Returns 1 on success and 0