Repository: randombit/botan
Pull request #3653 (2023-08-03 11:15:37)

Add PK_KEM_Decryptor::encapsulated_key_length
GH #3652


Patch:
diff --git a/src/lib/pubkey/kyber/kyber_common/kyber.cpp b/src/lib/pubkey/kyber/kyber_common/kyber.cpp
index dfa14daa0..320a0e1d1 100644
--- a/src/lib/pubkey/kyber/kyber_common/kyber.cpp
+++ b/src/lib/pubkey/kyber/kyber_common/kyber.cpp
@@ -1070,6 +1070,23 @@ class Kyber_KEM_Cryptor {
       const PolynomialMatrix m_at;
 };
 
+namespace {
+
+size_t kyber_key_length_to_encap_key_length(size_t kl) {
+   switch(kl) {
+      case 800:
+         return 768;
+      case 1184:
+         return 1088;
+      case 1568:
+         return 1568;
+      default:
+         throw Internal_Error("Unexpected Kyber key length");
+   }
+}
+
+}  // namespace
+
 class Kyber_KEM_Encryptor final : public PK_Ops::KEM_Encryption_with_KDF,
                                   protected Kyber_KEM_Cryptor {
    public:
@@ -1079,17 +1096,7 @@ class Kyber_KEM_Encryptor final : public PK_Ops::KEM_Encryption_with_KDF,
       size_t raw_kem_shared_key_length() const override { return 32; }
 
       size_t encapsulated_key_length() const override {
-         const size_t key_length = m_key.key_length();
-         switch(key_length) {
-            case 800:
-               return 768;
-            case 1184:
-               return 1088;
-            case 1568:
-               return 1568;
-            default:
-               throw Internal_Error("Unexpected Kyber key length");
-         }
+         return kyber_key_length_to_encap_key_length(m_key.key_length());
       }
 
       void raw_kem_encrypt(std::span<uint8_t> out_encapsulated_key,
@@ -1136,6 +1143,10 @@ class Kyber_KEM_Decryptor final : public PK_Ops::KEM_Decryption_with_KDF,
 
       size_t raw_kem_shared_key_length() const override { return 32; }
 
+      size_t encapsulated_key_length() const override {
+         return kyber_key_length_to_encap_key_length(m_key.key_length());
+      }
+
       void raw_kem_decrypt(std::span<uint8_t> out_shared_key, std::span<const uint8_t> encapsulated_key) override {
          // naming from kyber spec
          auto H = mode().H();
diff --git a/src/lib/pubkey/mce/mceliece_key.cpp b/src/lib/pubkey/mce/mceliece_key.cpp
index 359ac32fd..8e54fd562 100644
--- a/src/lib/pubkey/mce/mceliece_key.cpp
+++ b/src/lib/pubkey/mce/mceliece_key.cpp
@@ -323,6 +323,8 @@ class MCE_KEM_Decryptor final : public PK_Ops::KEM_Decryption_with_KDF {
          return ptext_sz + err_sz;
       }
 
+      size_t encapsulated_key_length() const override { return (m_key.get_code_length() + 7) / 8; }
+
       void raw_kem_decrypt(std::span<uint8_t> out_shared_key, std::span<const uint8_t> encapsulated_key) override {
          secure_vector<uint8_t> plaintext, error_mask;
          mceliece_decrypt(plaintext, error_mask, encapsulated_key.data(), encapsulated_key.size(), m_key);
diff --git a/src/lib/pubkey/pk_ops.h b/src/lib/pubkey/pk_ops.h
index 73aadc475..846ceb94b 100644
--- a/src/lib/pubkey/pk_ops.h
+++ b/src/lib/pubkey/pk_ops.h
@@ -161,6 +161,8 @@ class KEM_Decryption {
 
       virtual size_t shared_key_length(size_t desired_shared_key_len) const = 0;
 
+      virtual size_t encapsulated_key_length() const = 0;
+
       virtual ~KEM_Decryption() = default;
 };
 
diff --git a/src/lib/pubkey/pubkey.cpp b/src/lib/pubkey/pubkey.cpp
index e86d886a0..92a04e308 100644
--- a/src/lib/pubkey/pubkey.cpp
+++ b/src/lib/pubkey/pubkey.cpp
@@ -158,6 +158,10 @@ size_t PK_KEM_Decryptor::shared_key_length(size_t desired_shared_key_len) const
    return m_op->shared_key_length(desired_shared_key_len);
 }
 
+size_t PK_KEM_Decryptor::encapsulated_key_length() const {
+   return m_op->encapsulated_key_length();
+}
+
 PK_KEM_Decryptor::PK_KEM_Decryptor(const Private_Key& key,
                                    RandomNumberGenerator& rng,
                                    std::string_view param,
diff --git a/src/lib/pubkey/pubkey.h b/src/lib/pubkey/pubkey.h
index dedd2e269..123a852ba 100644
--- a/src/lib/pubkey/pubkey.h
+++ b/src/lib/pubkey/pubkey.h
@@ -755,6 +755,11 @@ class BOTAN_PUBLIC_API(2, 0) PK_KEM_Decryptor final {
       */
       size_t shared_key_length(size_t desired_shared_key_len) const;
 
+      /**
+      * Return the length of the encapsulated key expected by this KEM
+      */
+      size_t encapsulated_key_length() const;
+
       /**
       * Decrypts the shared key for data encryption.
       *
diff --git a/src/lib/pubkey/rsa/rsa.cpp b/src/lib/pubkey/rsa/rsa.cpp
index 441c5ddf4..44f6317a2 100644
--- a/src/lib/pubkey/rsa/rsa.cpp
+++ b/src/lib/pubkey/rsa/rsa.cpp
@@ -598,6 +598,8 @@ class RSA_KEM_Decryption_Operation final : public PK_Ops::KEM_Decryption_with_KD
 
       size_t raw_kem_shared_key_length() const override { return public_modulus_bytes(); }
 
+      size_t encapsulated_key_length() const override { return public_modulus_bytes(); }
+
       void raw_kem_decrypt(std::span<uint8_t> out_shared_key, std::span<const uint8_t> encapsulated_key) override {
          auto shared_key = raw_op(encapsulated_key.data(), encapsulated_key.size());
 


Test patch:
diff --git a/src/tests/test_pubkey.cpp b/src/tests/test_pubkey.cpp
index be2e82e60..0b293af43 100644
--- a/src/tests/test_pubkey.cpp
+++ b/src/tests/test_pubkey.cpp
@@ -463,6 +463,10 @@ Test::Result PK_KEM_Test::run_one_test(const std::string& /*header*/, const VarM
       return result;
    }
 
+   result.test_eq("encapsulated key length matches expected",
+                  kem_result.encapsulated_shared_key().size(),
+                  dec->encapsulated_key_length());
+
    const Botan::secure_vector<uint8_t> decr_shared_key =
       dec->decrypt(C0.data(), C0.size(), desired_key_len, salt.data(), salt.size());