Fix ML-DSA private key encoding

src/lib/pubkey/dilithium/dilithium_common/dilithium.cpp

@@ -420,13 +420,26 @@ Dilithium_PrivateKey::Dilithium_PrivateKey(std::span<const uint8_t> sk, Dilithiu
 }
 
 secure_vector<uint8_t> Dilithium_PrivateKey::raw_private_key_bits() const {
-   return this->private_key_bits();
+   const auto& seed_opt = this->m_private->seed();
+   if(!seed_opt.has_value()) {
+      throw Invalid_State(
+         "cannot return Dilithium or ML-DSA private key raw bits, since the key does not contain the seed");
+   }
+   return seed_opt.value().get();
 }
 
 secure_vector<uint8_t> Dilithium_PrivateKey::private_key_bits() const {
    return m_private->mode().keypair_codec().encode_keypair({m_public, m_private});
 }
 
+bool Dilithium_PrivateKey::is_mldsa() const {
+   return m_private->mode().is_ml_dsa();
+}
+
+bool Dilithium_PrivateKey::is_dilithium_round3() const {
+   return !m_private->mode().is_ml_dsa();
+}
+
 std::unique_ptr<PK_Ops::Signature> Dilithium_PrivateKey::create_signature_op(RandomNumberGenerator& rng,
                                                                              std::string_view params,
                                                                              std::string_view provider) const {

src/lib/pubkey/dilithium/dilithium_common/dilithium.h

@@ -142,6 +142,10 @@ class BOTAN_PUBLIC_API(3, 0) Dilithium_PrivateKey final : public virtual Dilithi
                                                              std::string_view params,
                                                              std::string_view provider) const override;
 
+      bool is_mldsa() const;
+
+      bool is_dilithium_round3() const;
+
    private:
       friend class Dilithium_Signature_Operation;
 

src/lib/pubkey/dilithium/dilithium_common/dilithium_algos.cpp

@@ -137,8 +137,6 @@ void poly_pack_gamma1(const DilithiumPoly& p, BufferStuffer& stuffer, const Dili
    BOTAN_ASSERT_UNREACHABLE();
 }
 
-#if defined(BOTAN_NEEDS_DILITHIUM_PRIVATE_KEY_ENCODING)
-
 /**
  * NIST FIPS 204, Algorithm 17 (BitPack)
  * (for a = -eta, b = eta)
@@ -164,8 +162,6 @@ void poly_pack_t0(const DilithiumPoly& p, BufferStuffer& stuffer) {
    poly_pack<TwoToTheDminus1 - 1, TwoToTheDminus1>(p, stuffer);
 }
 
-#endif
-
 /**
  * NIST FIPS 204, Algorithm 18 (SimpleBitUnpack)
  * (for a = 2^(bitlen(q-1)-d) - 1)
@@ -196,8 +192,6 @@ void poly_unpack_gamma1(DilithiumPoly& p, ByteSourceT& byte_source, const Dilith
    BOTAN_ASSERT_UNREACHABLE();
 }
 
-#if defined(BOTAN_NEEDS_DILITHIUM_PRIVATE_KEY_ENCODING)
-
 /**
  * NIST FIPS 204, Algorithm 19 (BitUnpack)
  * (for a = -eta, b = eta)
@@ -223,8 +217,6 @@ void poly_unpack_t0(DilithiumPoly& p, BufferSlicer& slicer) {
    poly_unpack<TwoToTheDminus1 - 1, TwoToTheDminus1>(p, slicer);
 }
 
-#endif
-
 /**
  * NIST FIPS 204, Algorithm 20 (HintBitPack)
  */
@@ -363,8 +355,6 @@ std::pair<DilithiumSeedRho, DilithiumPolyVec> decode_public_key(StrongSpan<const
    return {std::move(rho), std::move(t1)};
 }
 
-#if defined(BOTAN_NEEDS_DILITHIUM_PRIVATE_KEY_ENCODING)
-
 /**
  * NIST FIPS 204, Algorithm 24 (skEncode)
  */
@@ -469,8 +459,6 @@ DilithiumInternalKeypair decode_keypair(StrongSpan<const DilithiumSerializedPriv
    return keypair;
 }
 
-#endif
-
 /**
  * NIST FIPS 204, Algorithm 26 (sigEncode)
  */

src/lib/pubkey/dilithium/dilithium_common/dilithium_algos.h

@@ -14,12 +14,6 @@
 
 #include <botan/internal/dilithium_types.h>
 
-// ML-DSA does encode the private key only by its random seeds.
-#if defined(BOTAN_HAS_DILITHIUM) || defined(BOTAN_HAS_DILITHIUM_AES)
-   // NOLINTNEXTLINE(*-macro-usage)
-   #define BOTAN_NEEDS_DILITHIUM_PRIVATE_KEY_ENCODING 1
-#endif
-
 namespace Botan::Dilithium_Algos {
 
 DilithiumInternalKeypair expand_keypair(DilithiumSeedRandomness xi, DilithiumConstants mode);
@@ -51,15 +45,10 @@ DilithiumSerializedPublicKey encode_public_key(StrongSpan<const DilithiumSeedRho
 
 std::pair<DilithiumSeedRho, DilithiumPolyVec> decode_public_key(StrongSpan<const DilithiumSerializedPublicKey> pk,
                                                                 const DilithiumConstants& mode);
-
-#if BOTAN_NEEDS_DILITHIUM_PRIVATE_KEY_ENCODING
-
 DilithiumSerializedPrivateKey encode_keypair(const DilithiumInternalKeypair& keypair);
 
 DilithiumInternalKeypair decode_keypair(StrongSpan<const DilithiumSerializedPrivateKey> sk, DilithiumConstants mode);
 
-#endif
-
 std::pair<DilithiumPolyVec, DilithiumPolyVec> power2round(const DilithiumPolyVec& vec);
 
 std::pair<DilithiumPolyVec, DilithiumPolyVec> decompose(const DilithiumPolyVec& vec, const DilithiumConstants& mode);

src/lib/pubkey/dilithium/ml_dsa/ml_dsa_impl.cpp

@@ -8,20 +8,113 @@
 
 #include <botan/internal/ml_dsa_impl.h>
 
+#include <botan/asn1_obj.h>
+#include <botan/ber_dec.h>
+#include <botan/der_enc.h>
+#include <botan/exceptn.h>
 #include <botan/internal/dilithium_algos.h>
+#include <botan/internal/dilithium_types.h>
+#include <utility>
 
 namespace Botan {
 
+namespace {
+typedef Botan::DilithiumInternalKeypair (*decoding_func)(std::span<const uint8_t> key_bits,
+                                                         Botan::DilithiumConstants mode);
+
+Botan::DilithiumInternalKeypair decode_seed_only(std::span<const uint8_t> key_bits, Botan::DilithiumConstants mode) {
+   Botan::secure_vector<uint8_t> seed;
+   Botan::BER_Decoder(key_bits)
+      .decode(seed, Botan::ASN1_Type::OctetString, Botan::ASN1_Type(0), Botan::ASN1_Class::ContextSpecific)
+      .verify_end();
+   return Botan::Dilithium_Algos::expand_keypair(Botan::DilithiumSeedRandomness(seed), std::move(mode));
+}
+
+Botan::DilithiumInternalKeypair decode_expanded_only(std::span<const uint8_t> key_bits,
+                                                     Botan::DilithiumConstants mode) {
+   Botan::secure_vector<uint8_t> expanded;
+   Botan::BER_Decoder(key_bits).decode(expanded, Botan::ASN1_Type::OctetString).verify_end();
+   Botan::DilithiumInternalKeypair key_pair =
+      Botan::Dilithium_Algos::decode_keypair(Botan::DilithiumSerializedPrivateKey(expanded), std::move(mode));
+   return key_pair;
+}
+
+Botan::DilithiumInternalKeypair decode_seed_plus_expanded(std::span<const uint8_t> key_bits,
+                                                          Botan::DilithiumConstants mode) {
+   Botan::secure_vector<uint8_t> expanded;
+   Botan::secure_vector<uint8_t> seed;
+   Botan::BER_Decoder(key_bits)
+      .start_sequence()
+      .decode(seed, Botan::ASN1_Type::OctetString)
+      .decode(expanded, Botan::ASN1_Type::OctetString)
+      .end_cons()
+      .verify_end();
+   const Botan::DilithiumInternalKeypair key_pair =
+      Botan::Dilithium_Algos::decode_keypair(Botan::DilithiumSerializedPrivateKey(expanded), mode);
+   const Botan::DilithiumInternalKeypair key_pair_from_seed =
+      Botan::Dilithium_Algos::expand_keypair(Botan::DilithiumSeedRandomness(seed), std::move(mode));
+
+   DilithiumSerializedPrivateKey expanded_from_seed = Dilithium_Algos::encode_keypair(key_pair_from_seed);
+
+   if(expanded_from_seed.get() != expanded) {
+      throw Botan::Decoding_Error("seed and expanded key in ML-DSA serialized key do not match");
+   }
+   return key_pair_from_seed;
+}
+
+}  // namespace
+
 secure_vector<uint8_t> ML_DSA_Expanding_Keypair_Codec::encode_keypair(DilithiumInternalKeypair keypair) const {
    BOTAN_ASSERT_NONNULL(keypair.second);
    const auto& seed = keypair.second->seed();
-   BOTAN_ARG_CHECK(seed.has_value(), "Cannot encode keypair without the private seed");
-   return seed.value().get();
+   if(!keypair.second->mode().is_ml_dsa()) {
+      // return the raw seed for dilithium
+      BOTAN_ARG_CHECK(seed.has_value(), "Cannot encode keypair without the private seed");
+      return seed.value().get();
+   }
+   secure_vector<uint8_t> result;
+   DER_Encoder der_enc(result);
+   if(!seed.has_value()) {
+      // encode expanded-only format
+      DilithiumSerializedPrivateKey expanded = Dilithium_Algos::encode_keypair(keypair);
+      der_enc.encode(expanded.get(), ASN1_Type::OctetString);
+   } else {
+      // encode seed-only format
+      der_enc.encode(seed.value().get(),
+                     ASN1_Type(Botan::ASN1_Type::OctetString) /*real_type*/,
+                     ASN1_Type(Botan::ASN1_Type(0)),
+                     Botan::ASN1_Class::ContextSpecific);
+      /* 
+       * This yields the following ASN.1/DER structure:
+       <80 20>
+       0  32: [0]
+            :   00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F
+            :   10 11 12 13 14 15 16 17 18 19 1A 1B 1C 1D 1E 1F
+       *
+       *  Note: The previous format, which only contains the seed as an OCTET STRING without the [0]-tag, can still be decoded by OpenSSL. Apparently, it has this feature as a non-standard compatibility fallback for legacy formats.
+       */
+   }
+
+   return result;
 }
 
-DilithiumInternalKeypair ML_DSA_Expanding_Keypair_Codec::decode_keypair(std::span<const uint8_t> private_key_seed,
+DilithiumInternalKeypair ML_DSA_Expanding_Keypair_Codec::decode_keypair(std::span<const uint8_t> private_key_bits,
                                                                         DilithiumConstants mode) const {
-   return Dilithium_Algos::expand_keypair(DilithiumSeedRandomness(private_key_seed), std::move(mode));
+   if(private_key_bits.size() == 32) {
+      // backwards compatibility (not RFC 9881 conforming) to raw seed format.
+      return Botan::Dilithium_Algos::expand_keypair(Botan::DilithiumSeedRandomness(private_key_bits), std::move(mode));
+   }
+   // "seed-only" format from RFC 9881
+   BER_Decoder ber_dec(private_key_bits);
+   auto obj = ber_dec.peek_next_object();
+   if(obj.type() == ASN1_Type(0)) {
+      return decode_seed_only(private_key_bits, mode);
+   }
+   // now it could still be "expanded-only" or "both"
+   if(obj.type() == ASN1_Type::OctetString) {
+      return decode_expanded_only(private_key_bits, mode);
+   }
+   return decode_seed_plus_expanded(private_key_bits, mode);
 }
 
 }  // namespace Botan

src/tests/test_dilithium.cpp

@@ -9,16 +9,20 @@
  * Botan is released under the Simplified BSD License (see license.txt)
  */
 
+#include "test_rng.h"
 #include "tests.h"
 
+#include <memory>
+#include <vector>
+
 #if defined(BOTAN_HAS_DILITHIUM_COMMON)
    #include <botan/dilithium.h>
    #include <botan/hash.h>
    #include <botan/pk_algs.h>
+   #include <botan/pk_keys.h>
    #include <botan/pubkey.h>
 
    #include "test_pubkey.h"
-   #include "test_rng.h"
 #endif
 
 namespace Botan_Tests {
@@ -50,8 +54,10 @@ class Dilithium_KAT_Tests : public Text_Based_Test {
 
          const Botan::Dilithium_PrivateKey priv_key(*dilithium_test_rng, DerivedT::mode);
 
-         result.test_bin_eq(
-            "generated expected private key hash", sha3_256->process(priv_key.private_key_bits()), ref_sk_hash);
+         if(!priv_key.is_mldsa()) {
+            result.test_bin_eq(
+               "generated expected private key hash", sha3_256->process(priv_key.private_key_bits()), ref_sk_hash);
+         }
 
          result.test_bin_eq(
             "generated expected public key hash", sha3_256->process(priv_key.public_key_bits()), ref_pk_hash);
@@ -288,4 +294,56 @@ BOTAN_REGISTER_TEST("pubkey", "dilithium_keygen", Dilithium_Keygen_Tests);
 
 }  // namespace
 
+#if defined(BOTAN_HAS_DILITHIUM_COMMON) && defined(BOTAN_HAS_SHA3)
+class MLDSA_Privkey_Tests : public Text_Based_Test {
+   public:
+      MLDSA_Privkey_Tests() : Text_Based_Test("mldsa_privkey.vec", "key") {}
+
+      Test::Result run_one_test(const std::string& name, const VarMap& vars) override {
+         Test::Result result(name);
+         const std::vector<uint8_t> key_bits = vars.get_req_bin("key");
+         bool expect_decoding_failure = false;
+         Botan::DilithiumMode mode = Botan::DilithiumMode::ML_DSA_4x4;
+         if(name.starts_with("mldsa-44")) {
+            mode = Botan::DilithiumMode::ML_DSA_4x4;
+         } else if(name.starts_with("mldsa-65")) {
+            mode = Botan::DilithiumMode::ML_DSA_6x5;
+         } else if(name.starts_with("mldsa-87")) {
+            mode = Botan::DilithiumMode::ML_DSA_8x7;
+         } else {
+            throw Botan_Tests::Test_Error(
+               "internal error for ML-DSA test vector: encountered unknown ml-dsa mode string (test-case-specific token)");
+         }
+         if(name.ends_with("-invalid")) {
+            expect_decoding_failure = true;
+         }
+         std::unique_ptr<Botan::Dilithium_PrivateKey> priv_key;
+         try {
+            priv_key = std::make_unique<Botan::Dilithium_PrivateKey>(key_bits, mode);
+         } catch(const Botan::Decoding_Error&) {
+            result.test_is_true("invalid ML-DSA key rejected", expect_decoding_failure);
+            return result;
+         }
+         std::vector<uint8_t> ref_msg = {0, 1, 2, 4};
+         std::vector<uint8_t> rng_seed(48);
+         Botan_Tests::CTR_DRBG_AES256 rng(rng_seed);
+         auto signer = Botan::PK_Signer(*priv_key, rng, "Randomized");
+         auto signature = signer.sign_message(ref_msg.data(), ref_msg.size(), rng);
+
+         const Botan::Dilithium_PublicKey pub_key(priv_key->public_key_bits(), mode);
+         auto verifier = Botan::PK_Verifier(pub_key, "");
+         verifier.update(ref_msg.data(), ref_msg.size());
+         result.test_is_true("signature verifies", verifier.check_signature(signature.data(), signature.size()));
+
+         auto reencoded_priv_key = priv_key->private_key_bits();
+         auto redecoded_priv_key = Botan::Dilithium_PrivateKey(reencoded_priv_key, mode);
+         result.test_is_true("re-encoding and subsequent re-decoding of private ML-DSA key without error", true);
+         return result;
+      }
+};
+
+BOTAN_REGISTER_TEST("pubkey", "mldsa_private_key", MLDSA_Privkey_Tests);
+
+#endif
+
 }  // namespace Botan_Tests