Made Serializable::write_exact for allocation-free serialization (#48)

* Removed useless use statements

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Pending
+
+* Added `Serializable::write_exact` so serialization requires less stack space
+
 ## [0.11.0] - 2023-10-11
 
 ### Removals

src/aead.rs

@@ -4,7 +4,7 @@ use crate::{
     kdf::{Kdf as KdfTrait, LabeledExpand, SimpleHkdf},
     kem::Kem as KemTrait,
     setup::ExporterSecret,
-    util::{enforce_equal_len, full_suite_id, FullSuiteId},
+    util::{enforce_equal_len, enforce_outbuf_len, full_suite_id, FullSuiteId},
     Deserializable, HpkeError, Serializable,
 };
 
@@ -139,8 +139,12 @@ impl<A: Aead> Default for AeadTag<A> {
 impl<A: Aead> Serializable for AeadTag<A> {
     type OutputSize = <A::AeadImpl as BaseAeadCore>::TagSize;
 
-    fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
-        self.0.clone()
+    // Pass to underlying to_bytes() impl
+    fn write_exact(&self, buf: &mut [u8]) {
+        // Check the length is correct and panic if not
+        enforce_outbuf_len::<Self>(buf);
+
+        buf.copy_from_slice(&self.0);
     }
 }
 
@@ -760,4 +764,16 @@ mod test {
             crate::kem::DhP384HkdfSha384
         );
     }
+
+    /// Tests that Serialize::write_exact() panics when given a buffer of incorrect length
+    #[should_panic]
+    #[test]
+    fn test_write_exact() {
+        // Make an AES-GCM-128 tag (16 bytes) and try to serialize it to a buffer of 17 bytes. It
+        // shouldn't matter that this is sufficient room, since write_exact needs exactly the write
+        // size buffer
+        let tag = AeadTag::<AesGcm128>::default();
+        let mut buf = [0u8; 17];
+        tag.write_exact(&mut buf);
+    }
 }

src/dhkex.rs

@@ -74,12 +74,6 @@ pub trait DhKeyExchange {
 
 #[cfg(any(feature = "p256", feature = "p384"))]
 pub(crate) mod ecdh_nistp;
-#[cfg(feature = "p256")]
-pub use ecdh_nistp::p256::DhP256;
-#[cfg(feature = "p384")]
-pub use ecdh_nistp::p384::DhP384;
 
 #[cfg(feature = "x25519")]
 pub(crate) mod x25519;
-#[cfg(feature = "x25519")]
-pub use x25519::X25519;

src/dhkex/ecdh_nistp.rs

@@ -15,7 +15,7 @@ macro_rules! nistp_dhkex {
             use crate::{
                 dhkex::{DhError, DhKeyExchange},
                 kdf::{labeled_extract, Kdf as KdfTrait, LabeledExpand},
-                util::{enforce_equal_len, KemSuiteId},
+                util::{enforce_equal_len, enforce_outbuf_len, KemSuiteId},
                 Deserializable, HpkeError, Serializable,
             };
 
@@ -56,12 +56,16 @@ macro_rules! nistp_dhkex {
             impl Serializable for PublicKey {
                 type OutputSize = $pubkey_size;
 
-                fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
+                fn write_exact(&self, buf: &mut [u8]) {
+                    // Check the length is correct and panic if not
+                    enforce_outbuf_len::<Self>(buf);
+
                     // Get the uncompressed pubkey encoding
                     let encoded = self.0.as_affine().to_encoded_point(false);
                     // Serialize it
-                    GenericArray::clone_from_slice(encoded.as_bytes())
+                    buf.copy_from_slice(encoded.as_bytes());
                 }
+
             }
 
             // Everything is serialized and deserialized in uncompressed form
@@ -85,9 +89,12 @@ macro_rules! nistp_dhkex {
             impl Serializable for PrivateKey {
                 type OutputSize = $privkey_size;
 
-                fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
+                fn write_exact(&self, buf: &mut [u8]) {
+                    // Check the length is correct and panic if not
+                    enforce_outbuf_len::<Self>(buf);
+
                     // SecretKeys already know how to convert to bytes
-                    self.0.to_bytes()
+                    buf.copy_from_slice(&self.0.to_bytes());
                 }
             }
 
@@ -116,10 +123,13 @@ macro_rules! nistp_dhkex {
                 // resulting elliptic curve point.
                 type OutputSize = $ss_size;
 
-                fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
+                fn write_exact(&self, buf: &mut [u8]) {
+                    // Check the length is correct and panic if not
+                    enforce_outbuf_len::<Self>(buf);
+
                     // elliptic_curve::ecdh::SharedSecret::raw_secret_bytes returns the serialized
                     // x-coordinate
-                    *self.0.raw_secret_bytes()
+                    buf.copy_from_slice(self.0.raw_secret_bytes())
                 }
             }
 

src/dhkex/x25519.rs

@@ -1,14 +1,11 @@
 use crate::{
     dhkex::{DhError, DhKeyExchange},
     kdf::{labeled_extract, Kdf as KdfTrait, LabeledExpand},
-    util::{enforce_equal_len, KemSuiteId},
+    util::{enforce_equal_len, enforce_outbuf_len, KemSuiteId},
     Deserializable, HpkeError, Serializable,
 };
 
-use generic_array::{
-    typenum::{self, Unsigned},
-    GenericArray,
-};
+use generic_array::typenum::{self, Unsigned};
 use subtle::{Choice, ConstantTimeEq};
 
 // We wrap the types in order to abstract away the dalek dep
@@ -46,8 +43,11 @@ impl Serializable for PublicKey {
     type OutputSize = typenum::U32;
 
     // Dalek lets us convert pubkeys to [u8; 32]
-    fn to_bytes(&self) -> GenericArray<u8, typenum::U32> {
-        GenericArray::clone_from_slice(self.0.as_bytes())
+    fn write_exact(&self, buf: &mut [u8]) {
+        // Check the length is correct and panic if not
+        enforce_outbuf_len::<Self>(buf);
+
+        buf.copy_from_slice(self.0.as_bytes());
     }
 }
 
@@ -70,8 +70,11 @@ impl Serializable for PrivateKey {
     type OutputSize = typenum::U32;
 
     // Dalek lets us convert scalars to [u8; 32]
-    fn to_bytes(&self) -> GenericArray<u8, typenum::U32> {
-        GenericArray::clone_from_slice(&self.0.to_bytes())
+    fn write_exact(&self, buf: &mut [u8]) {
+        // Check the length is correct and panic if not
+        enforce_outbuf_len::<Self>(buf);
+
+        buf.copy_from_slice(self.0.as_bytes());
     }
 }
 impl Deserializable for PrivateKey {
@@ -102,9 +105,12 @@ impl Serializable for KexResult {
     type OutputSize = typenum::U32;
 
     // curve25519's point representation is our DH result. We don't have to do anything special.
-    fn to_bytes(&self) -> GenericArray<u8, typenum::U32> {
+    fn write_exact(&self, buf: &mut [u8]) {
+        // Check the length is correct and panic if not
+        enforce_outbuf_len::<Self>(buf);
+
         // Dalek lets us convert shared secrets to to [u8; 32]
-        GenericArray::clone_from_slice(self.0.as_bytes())
+        buf.copy_from_slice(self.0.as_bytes());
     }
 }
 

src/kem/dhkem.rs

@@ -15,12 +15,11 @@ macro_rules! impl_dhkem {
                 dhkex::{DhKeyExchange, MAX_PUBKEY_SIZE},
                 kdf::{extract_and_expand, Kdf as KdfTrait},
                 kem::{Kem as KemTrait, SharedSecret},
-                util::kem_suite_id,
+                util::{enforce_outbuf_len, kem_suite_id},
                 Deserializable, HpkeError, Serializable,
             };
 
             use digest::OutputSizeUser;
-            use generic_array::GenericArray;
             use rand_core::{CryptoRng, RngCore};
 
             // Define convenience types
@@ -46,8 +45,11 @@ macro_rules! impl_dhkem {
                 type OutputSize = <<$dhkex as DhKeyExchange>::PublicKey as Serializable>::OutputSize;
 
                 // Pass to underlying to_bytes() impl
-                fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
-                    self.0.to_bytes()
+                fn write_exact(&self, buf: &mut [u8]) {
+                    // Check the length is correct and panic if not
+                    enforce_outbuf_len::<Self>(buf);
+
+                    buf.copy_from_slice(&self.0.to_bytes());
                 }
             }
 

src/lib.rs

@@ -189,9 +189,24 @@ impl core::fmt::Display for HpkeError {
 
 /// Implemented by types that have a fixed-length byte representation
 pub trait Serializable {
+    /// Serialized size in bytes
     type OutputSize: ArrayLength<u8>;
 
-    fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize>;
+    /// Serializes `self` to the given slice. `buf` MUST have length equal to `Self::size()`.
+    ///
+    /// Panics
+    /// ======
+    /// Panics if `buf.len() != Self::size()`.
+    fn write_exact(&self, buf: &mut [u8]);
+
+    /// Serializes `self` to a new array
+    fn to_bytes(&self) -> GenericArray<u8, Self::OutputSize> {
+        // Make a buffer of the correct size and write to it
+        let mut buf = GenericArray::default();
+        self.write_exact(&mut buf);
+        // Return the buffer
+        buf
+    }
 
     /// Returns the size (in bytes) of this type when serialized
     fn size() -> usize {

src/util.rs

@@ -1,4 +1,4 @@
-use crate::{aead::Aead, kdf::Kdf as KdfTrait, kem::Kem as KemTrait, HpkeError};
+use crate::{aead::Aead, kdf::Kdf as KdfTrait, kem::Kem as KemTrait, HpkeError, Serializable};
 
 use byteorder::{BigEndian, ByteOrder};
 
@@ -93,3 +93,16 @@ pub(crate) fn enforce_equal_len(expected_len: usize, given_len: usize) -> Result
         Ok(())
     }
 }
+
+/// Helper function for `Serializable::write_exact`. Takes a buffer and a serializable type `T` and
+/// panics iff `buf.len() != T::size()`.
+pub(crate) fn enforce_outbuf_len<T: Serializable>(buf: &[u8]) {
+    let size = T::size();
+    let buf_len = buf.len();
+    assert!(
+        size == buf_len,
+        "write_exact(): serialized size ({}) does not equal buffer length ({})",
+        size,
+        buf_len,
+    );
+}