From c4b3341ec478fc9632032d4718a27db2f60f94af Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 9 Nov 2023 12:58:06 +1100
Subject: [PATCH 01/20] Add benchmark

---
 Cargo.lock                                    |  1 +
 misc/quick-protobuf-codec/Cargo.toml          |  7 +++
 misc/quick-protobuf-codec/benches/codec.rs    | 28 +++++++++++
 .../quick-protobuf-codec/src/generated/mod.rs |  2 +
 .../src/generated/test.proto                  |  7 +++
 .../src/generated/test.rs                     | 47 +++++++++++++++++++
 misc/quick-protobuf-codec/src/lib.rs          |  5 ++
 7 files changed, 97 insertions(+)
 create mode 100644 misc/quick-protobuf-codec/benches/codec.rs
 create mode 100644 misc/quick-protobuf-codec/src/generated/mod.rs
 create mode 100644 misc/quick-protobuf-codec/src/generated/test.proto
 create mode 100644 misc/quick-protobuf-codec/src/generated/test.rs

diff --git a/Cargo.lock b/Cargo.lock
index 8c252a1d8b8..57c40cb087b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4370,6 +4370,7 @@ version = "0.2.0"
 dependencies = [
  "asynchronous-codec",
  "bytes",
+ "criterion",
  "quick-protobuf",
  "thiserror",
  "unsigned-varint 0.8.0",
diff --git a/misc/quick-protobuf-codec/Cargo.toml b/misc/quick-protobuf-codec/Cargo.toml
index 2e309a02889..e456252cb5a 100644
--- a/misc/quick-protobuf-codec/Cargo.toml
+++ b/misc/quick-protobuf-codec/Cargo.toml
@@ -17,6 +17,13 @@ thiserror = "1.0"
 unsigned-varint = { workspace = true, features = ["asynchronous_codec"] }
 quick-protobuf = "0.8"
 
+[dev-dependencies]
+criterion = "0.5.1"
+
+[[bench]]
+name = "codec"
+harness = false
+
 # Passing arguments to the docsrs builder in order to properly document cfg's.
 # More information: https://docs.rs/about/builds#cross-compiling
 [package.metadata.docs.rs]
diff --git a/misc/quick-protobuf-codec/benches/codec.rs b/misc/quick-protobuf-codec/benches/codec.rs
new file mode 100644
index 00000000000..0f6ce9469c5
--- /dev/null
+++ b/misc/quick-protobuf-codec/benches/codec.rs
@@ -0,0 +1,28 @@
+use asynchronous_codec::Encoder;
+use bytes::BytesMut;
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use quick_protobuf_codec::{proto, Codec};
+
+pub fn benchmark(c: &mut Criterion) {
+    for size in [1000, 10_000, 100_000, 1_000_000, 10_000_000] {
+        c.bench_with_input(BenchmarkId::new("encode", size), &size, |b, i| {
+            b.iter_batched(
+                || {
+                    let mut out = BytesMut::new();
+                    out.reserve(i + 100);
+                    let codec = Codec::<proto::Message>::new(i + 100);
+                    let msg = proto::Message {
+                        data: vec![0; size],
+                    };
+
+                    (codec, out, msg)
+                },
+                |(mut codec, mut out, msg)| codec.encode(msg, &mut out).unwrap(),
+                BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
+criterion_group!(benches, benchmark);
+criterion_main!(benches);
diff --git a/misc/quick-protobuf-codec/src/generated/mod.rs b/misc/quick-protobuf-codec/src/generated/mod.rs
new file mode 100644
index 00000000000..b9f982f8dfd
--- /dev/null
+++ b/misc/quick-protobuf-codec/src/generated/mod.rs
@@ -0,0 +1,2 @@
+// Automatically generated mod.rs
+pub mod test;
diff --git a/misc/quick-protobuf-codec/src/generated/test.proto b/misc/quick-protobuf-codec/src/generated/test.proto
new file mode 100644
index 00000000000..5b1f46c0bfa
--- /dev/null
+++ b/misc/quick-protobuf-codec/src/generated/test.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test;
+
+message Message {
+    bytes data = 1;
+}
diff --git a/misc/quick-protobuf-codec/src/generated/test.rs b/misc/quick-protobuf-codec/src/generated/test.rs
new file mode 100644
index 00000000000..b353e6d9183
--- /dev/null
+++ b/misc/quick-protobuf-codec/src/generated/test.rs
@@ -0,0 +1,47 @@
+// Automatically generated rust module for 'test.proto' file
+
+#![allow(non_snake_case)]
+#![allow(non_upper_case_globals)]
+#![allow(non_camel_case_types)]
+#![allow(unused_imports)]
+#![allow(unknown_lints)]
+#![allow(clippy::all)]
+#![cfg_attr(rustfmt, rustfmt_skip)]
+
+
+use quick_protobuf::{MessageInfo, MessageRead, MessageWrite, BytesReader, Writer, WriterBackend, Result};
+use quick_protobuf::sizeofs::*;
+use super::*;
+
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Debug, Default, PartialEq, Clone)]
+pub struct Message {
+    pub data: Vec<u8>,
+}
+
+impl<'a> MessageRead<'a> for Message {
+    fn from_reader(r: &mut BytesReader, bytes: &'a [u8]) -> Result<Self> {
+        let mut msg = Self::default();
+        while !r.is_eof() {
+            match r.next_tag(bytes) {
+                Ok(10) => msg.data = r.read_bytes(bytes)?.to_owned(),
+                Ok(t) => { r.read_unknown(bytes, t)?; }
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(msg)
+    }
+}
+
+impl MessageWrite for Message {
+    fn get_size(&self) -> usize {
+        0
+        + if self.data.is_empty() { 0 } else { 1 + sizeof_len((&self.data).len()) }
+    }
+
+    fn write_message<W: WriterBackend>(&self, w: &mut Writer<W>) -> Result<()> {
+        if !self.data.is_empty() { w.write_with_tag(10, |w| w.write_bytes(&**&self.data))?; }
+        Ok(())
+    }
+}
+
diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 2d1fda99a70..bf0d36cc83c 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -6,6 +6,11 @@ use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer};
 use std::marker::PhantomData;
 use unsigned_varint::codec::UviBytes;
 
+mod generated;
+
+#[doc(hidden)] // NOT public API. Do not use.
+pub use generated::test as proto;
+
 /// [`Codec`] implements [`Encoder`] and [`Decoder`], uses [`unsigned_varint`]
 /// to prefix messages with their length and uses [`quick_protobuf`] and a provided
 /// `struct` implementing [`MessageRead`] and [`MessageWrite`] to do the encoding.

From ef00f01a0b23423941f42476b43ca43dd80ffb2c Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 9 Nov 2023 13:32:17 +1100
Subject: [PATCH 02/20] Add test for encoding large message

---
 misc/quick-protobuf-codec/tests/large_message.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)
 create mode 100644 misc/quick-protobuf-codec/tests/large_message.rs

diff --git a/misc/quick-protobuf-codec/tests/large_message.rs b/misc/quick-protobuf-codec/tests/large_message.rs
new file mode 100644
index 00000000000..65dafe065d1
--- /dev/null
+++ b/misc/quick-protobuf-codec/tests/large_message.rs
@@ -0,0 +1,16 @@
+use asynchronous_codec::Encoder;
+use bytes::BytesMut;
+use quick_protobuf_codec::proto;
+use quick_protobuf_codec::Codec;
+
+#[test]
+fn encode_large_message() {
+    let mut codec = Codec::<proto::Message>::new(1_001_000);
+    let mut dst = BytesMut::new();
+    dst.reserve(1_001_000);
+    let message = proto::Message {
+        data: vec![0; 1_000_000],
+    };
+
+    codec.encode(message, &mut dst).unwrap();
+}

From 8f3df3961ac9d8aa1c4ae9ce0d9f5b3c2132c7fd Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 12:34:30 +1100
Subject: [PATCH 03/20] Don't allocate as much

---
 misc/quick-protobuf-codec/Cargo.toml |  2 +-
 misc/quick-protobuf-codec/src/lib.rs | 45 ++++++++++++++++++++--------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/misc/quick-protobuf-codec/Cargo.toml b/misc/quick-protobuf-codec/Cargo.toml
index e456252cb5a..133bbd2bb6e 100644
--- a/misc/quick-protobuf-codec/Cargo.toml
+++ b/misc/quick-protobuf-codec/Cargo.toml
@@ -14,7 +14,7 @@ categories = ["asynchronous"]
 asynchronous-codec = { workspace = true }
 bytes = { version = "1" }
 thiserror = "1.0"
-unsigned-varint = { workspace = true, features = ["asynchronous_codec"] }
+unsigned-varint = { workspace = true, features = ["std"] }
 quick-protobuf = "0.8"
 
 [dev-dependencies]
diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index bf0d36cc83c..3d52229c32e 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -1,10 +1,10 @@
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
 use asynchronous_codec::{Decoder, Encoder};
-use bytes::{Bytes, BytesMut};
+use bytes::{Buf, BytesMut};
 use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer};
+use std::io;
 use std::marker::PhantomData;
-use unsigned_varint::codec::UviBytes;
 
 mod generated;
 
@@ -15,7 +15,7 @@ pub use generated::test as proto;
 /// to prefix messages with their length and uses [`quick_protobuf`] and a provided
 /// `struct` implementing [`MessageRead`] and [`MessageWrite`] to do the encoding.
 pub struct Codec<In, Out = In> {
-    uvi: UviBytes,
+    max_message_len_bytes: usize,
     phantom: PhantomData<(In, Out)>,
 }
 
@@ -26,10 +26,8 @@ impl<In, Out> Codec<In, Out> {
     /// Protobuf message. The limit does not include the bytes needed for the
     /// [`unsigned_varint`].
     pub fn new(max_message_len_bytes: usize) -> Self {
-        let mut uvi = UviBytes::default();
-        uvi.set_max_len(max_message_len_bytes);
         Self {
-            uvi,
+            max_message_len_bytes,
             phantom: PhantomData,
         }
     }
@@ -40,11 +38,16 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
     type Error = Error;
 
     fn encode(&mut self, item: Self::Item<'_>, dst: &mut BytesMut) -> Result<(), Self::Error> {
-        let mut encoded_msg = Vec::new();
-        let mut writer = Writer::new(&mut encoded_msg);
+        let message_length = item.get_size();
+
+        let mut uvi_buf = unsigned_varint::encode::usize_buffer();
+        let encoded_length = unsigned_varint::encode::usize(message_length, &mut uvi_buf);
+
+        dst.extend_from_slice(encoded_length);
+
+        let mut writer = Writer::new(dst.as_mut());
         item.write_message(&mut writer)
             .expect("Encoding to succeed");
-        self.uvi.encode(Bytes::from(encoded_msg), dst)?;
 
         Ok(())
     }
@@ -58,14 +61,30 @@ where
     type Error = Error;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
-        let msg = match self.uvi.decode(src)? {
-            None => return Ok(None),
-            Some(msg) => msg,
+        let (len, remaining) = match unsigned_varint::decode::usize(src) {
+            Ok((len, remaining)) => (len, remaining),
+            Err(unsigned_varint::decode::Error::Insufficient) => return Ok(None),
+            Err(e) => return Err(Error(io::Error::new(io::ErrorKind::InvalidData, e))),
         };
+        let consumed = src.len() - remaining.len();
+        src.advance(consumed);
+
+        if len > self.max_message_len_bytes {
+            return Err(Error(io::Error::new(
+                io::ErrorKind::PermissionDenied,
+                format!(
+                    "message with {len}b exceeds maximum of {}b",
+                    self.max_message_len_bytes
+                ),
+            )));
+        }
+
+        let msg = src.split_to(len);
 
         let mut reader = BytesReader::from_bytes(&msg);
         let message = Self::Item::from_reader(&mut reader, &msg)
-            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
         Ok(Some(message))
     }
 }

From 58a27fbc22a764b5c04b0a66dab5757b06c1f018 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 12:36:06 +1100
Subject: [PATCH 04/20] Add changelog entry

---
 Cargo.lock                             | 2 +-
 Cargo.toml                             | 2 +-
 misc/quick-protobuf-codec/CHANGELOG.md | 5 +++++
 misc/quick-protobuf-codec/Cargo.toml   | 2 +-
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 57c40cb087b..31e191996d9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4366,7 +4366,7 @@ dependencies = [
 
 [[package]]
 name = "quick-protobuf-codec"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
  "asynchronous-codec",
  "bytes",
diff --git a/Cargo.toml b/Cargo.toml
index a79c55bbf91..4daa24bcaa2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -117,7 +117,7 @@ multiaddr = "0.18.1"
 multihash = "0.19.1"
 multistream-select = { version = "0.13.0", path = "misc/multistream-select" }
 prometheus-client = "0.22.0"
-quick-protobuf-codec = { version = "0.2.0", path = "misc/quick-protobuf-codec" }
+quick-protobuf-codec = { version = "0.2.1", path = "misc/quick-protobuf-codec" }
 quickcheck = { package = "quickcheck-ext", path = "misc/quickcheck-ext" }
 rw-stream-sink = { version = "0.4.0", path = "misc/rw-stream-sink" }
 unsigned-varint = { version = "0.8.0" }
diff --git a/misc/quick-protobuf-codec/CHANGELOG.md b/misc/quick-protobuf-codec/CHANGELOG.md
index 740201f80d7..af36b1f0681 100644
--- a/misc/quick-protobuf-codec/CHANGELOG.md
+++ b/misc/quick-protobuf-codec/CHANGELOG.md
@@ -1,3 +1,8 @@
+## 0.2.1 - unreleased
+
+- Reduce allocations during encoding.
+  See [PR XXXX](https://github.com/libp2p/rust-libp2p/pull/XXXX).
+
 ## 0.2.0 
 
 - Raise MSRV to 1.65.
diff --git a/misc/quick-protobuf-codec/Cargo.toml b/misc/quick-protobuf-codec/Cargo.toml
index 133bbd2bb6e..1db8d8a2d0f 100644
--- a/misc/quick-protobuf-codec/Cargo.toml
+++ b/misc/quick-protobuf-codec/Cargo.toml
@@ -3,7 +3,7 @@ name = "quick-protobuf-codec"
 edition = "2021"
 rust-version = { workspace = true }
 description = "Asynchronous de-/encoding of Protobuf structs using asynchronous-codec, unsigned-varint and quick-protobuf."
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Max Inden <mail@max-inden.de>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"

From 466d049f298d632f05a2fa28fa5ce57b1b0cd666 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 12:54:25 +1100
Subject: [PATCH 05/20] Add test for max length

---
 Cargo.lock                           |  1 +
 misc/quick-protobuf-codec/Cargo.toml |  1 +
 misc/quick-protobuf-codec/src/lib.rs | 47 ++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 31e191996d9..f6efc970bc2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4371,6 +4371,7 @@ dependencies = [
  "asynchronous-codec",
  "bytes",
  "criterion",
+ "futures",
  "quick-protobuf",
  "thiserror",
  "unsigned-varint 0.8.0",
diff --git a/misc/quick-protobuf-codec/Cargo.toml b/misc/quick-protobuf-codec/Cargo.toml
index 1db8d8a2d0f..9c8ed89c231 100644
--- a/misc/quick-protobuf-codec/Cargo.toml
+++ b/misc/quick-protobuf-codec/Cargo.toml
@@ -19,6 +19,7 @@ quick-protobuf = "0.8"
 
 [dev-dependencies]
 criterion = "0.5.1"
+futures = "0.3.28"
 
 [[bench]]
 name = "codec"
diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 3d52229c32e..a2dd6f61375 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -91,10 +91,53 @@ where
 
 #[derive(thiserror::Error, Debug)]
 #[error("Failed to encode/decode message")]
-pub struct Error(#[from] std::io::Error);
+pub struct Error(#[from] io::Error);
 
-impl From<Error> for std::io::Error {
+impl From<Error> for io::Error {
     fn from(e: Error) -> Self {
         e.0
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use asynchronous_codec::FramedRead;
+    use futures::io::Cursor;
+    use futures::{FutureExt, StreamExt};
+    use std::error::Error;
+
+    #[test]
+    fn honors_max_message_length() {
+        let codec = Codec::<Dummy>::new(1);
+        let mut src = varint_zeroes(100);
+
+        let mut read = FramedRead::new(Cursor::new(&mut src), codec);
+        let err = read.next().now_or_never().unwrap().unwrap().unwrap_err();
+
+        assert_eq!(
+            err.source().unwrap().to_string(),
+            "message with 100b exceeds maximum of 1b"
+        )
+    }
+
+    /// Constructs a [`BytesMut`] of the provided length where the message is all zeros.
+    fn varint_zeroes(length: usize) -> BytesMut {
+        let mut buf = unsigned_varint::encode::usize_buffer();
+        let encoded_length = unsigned_varint::encode::usize(length, &mut buf);
+
+        let mut src = BytesMut::new();
+        src.extend_from_slice(encoded_length);
+        src.extend(std::iter::repeat(0).take(length));
+        src
+    }
+
+    #[derive(Debug)]
+    struct Dummy;
+
+    impl<'a> MessageRead<'a> for Dummy {
+        fn from_reader(_: &mut BytesReader, _: &'a [u8]) -> quick_protobuf::Result<Self> {
+            todo!()
+        }
+    }
+}

From 5b79d2d3b2dcba8a4486501d2895ece05cfd148f Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 13:15:50 +1100
Subject: [PATCH 06/20] Ensure we reserve enough space for new message

---
 misc/quick-protobuf-codec/src/lib.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index a2dd6f61375..f8f60694945 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -43,6 +43,7 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
         let mut uvi_buf = unsigned_varint::encode::usize_buffer();
         let encoded_length = unsigned_varint::encode::usize(message_length, &mut uvi_buf);
 
+        dst.reserve(message_length);
         dst.extend_from_slice(encoded_length);
 
         let mut writer = Writer::new(dst.as_mut());

From a43c1de19d71ffbd024843a89bb94ed6c99f4212 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 12:56:45 +1100
Subject: [PATCH 07/20] Update misc/quick-protobuf-codec/CHANGELOG.md

---
 misc/quick-protobuf-codec/CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/misc/quick-protobuf-codec/CHANGELOG.md b/misc/quick-protobuf-codec/CHANGELOG.md
index af36b1f0681..06abad13bce 100644
--- a/misc/quick-protobuf-codec/CHANGELOG.md
+++ b/misc/quick-protobuf-codec/CHANGELOG.md
@@ -1,7 +1,7 @@
 ## 0.2.1 - unreleased
 
 - Reduce allocations during encoding.
-  See [PR XXXX](https://github.com/libp2p/rust-libp2p/pull/XXXX).
+  See [PR 4782](https://github.com/libp2p/rust-libp2p/pull/4782).
 
 ## 0.2.0 
 

From 9959d16701d7d14f0dd7d8bd59855c0c25b59162 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:02:34 +1100
Subject: [PATCH 08/20] Ensure we have enough capacity in `BytesMut`

---
 misc/quick-protobuf-codec/src/lib.rs | 85 ++++++++++++++++++++++++++--
 protocols/gossipsub/src/protocol.rs  |  6 +-
 2 files changed, 84 insertions(+), 7 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index f8f60694945..98a16027af6 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -2,9 +2,10 @@
 
 use asynchronous_codec::{Decoder, Encoder};
 use bytes::{Buf, BytesMut};
-use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer};
+use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer, WriterBackend};
 use std::io;
 use std::marker::PhantomData;
+use std::mem::MaybeUninit;
 
 mod generated;
 
@@ -43,12 +44,88 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
         let mut uvi_buf = unsigned_varint::encode::usize_buffer();
         let encoded_length = unsigned_varint::encode::usize(message_length, &mut uvi_buf);
 
-        dst.reserve(message_length);
+        // Append the 'unsigned varint'-encoded length.
         dst.extend_from_slice(encoded_length);
 
-        let mut writer = Writer::new(dst.as_mut());
+        // Ensure we have enough capacity to encode our message.
+        dst.reserve(message_length);
+
+        let mut written = 0;
+
+        let mut writer = Writer::new(UninitMemoryWriterBackend::new(
+            dst.spare_capacity_mut(),
+            &mut written,
+        ));
         item.write_message(&mut writer)
-            .expect("Encoding to succeed");
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+
+        if written != message_length {
+            return Err(Error(io::Error::new(
+                io::ErrorKind::Other,
+                format!(
+                    "expected message to be {message_length} bytes long but was {written} bytes"
+                ),
+            )));
+        }
+
+        // SAFETY: `written` records exactly how many bytes we wrote to `dst`, thus it is safe to extend the length by `written`.
+        unsafe {
+            dst.set_len(dst.len() + written);
+        }
+
+        Ok(())
+    }
+}
+
+struct UninitMemoryWriterBackend<'a> {
+    memory: &'a mut [MaybeUninit<u8>],
+    written: &'a mut usize,
+}
+
+impl<'a> UninitMemoryWriterBackend<'a> {
+    fn new(memory: &'a mut [MaybeUninit<u8>], written: &'a mut usize) -> Self {
+        Self { memory, written }
+    }
+}
+
+impl<'a> WriterBackend for UninitMemoryWriterBackend<'a> {
+    fn pb_write_u8(&mut self, x: u8) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&[x])
+    }
+
+    fn pb_write_u32(&mut self, x: u32) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_i32(&mut self, x: i32) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_f32(&mut self, x: f32) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_u64(&mut self, x: u64) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_i64(&mut self, x: i64) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_f64(&mut self, x: f64) -> quick_protobuf::Result<()> {
+        self.pb_write_all(&x.to_le_bytes())
+    }
+
+    fn pb_write_all(&mut self, buf: &[u8]) -> quick_protobuf::Result<()> {
+        if self.memory.len() - *self.written < buf.len() {
+            return Err(quick_protobuf::errors::Error::UnexpectedEndOfBuffer);
+        }
+
+        for b in buf {
+            self.memory[*self.written].write(*b);
+            *self.written += 1;
+        }
 
         Ok(())
     }
diff --git a/protocols/gossipsub/src/protocol.rs b/protocols/gossipsub/src/protocol.rs
index 42d43c97510..3811055711d 100644
--- a/protocols/gossipsub/src/protocol.rs
+++ b/protocols/gossipsub/src/protocol.rs
@@ -573,7 +573,7 @@ mod tests {
     #[test]
     /// Test that RPC messages can be encoded and decoded successfully.
     fn encode_decode() {
-        fn prop(message: Message) {
+        fn prop(message: Message, buf_length: u32) {
             let message = message.0;
 
             let rpc = Rpc {
@@ -583,7 +583,7 @@ mod tests {
             };
 
             let mut codec = GossipsubCodec::new(u32::MAX as usize, ValidationMode::Strict);
-            let mut buf = BytesMut::new();
+            let mut buf = BytesMut::with_capacity(buf_length as usize);
             codec.encode(rpc.into_protobuf(), &mut buf).unwrap();
             let decoded_rpc = codec.decode(&mut buf).unwrap().unwrap();
             // mark as validated as its a published message
@@ -597,7 +597,7 @@ mod tests {
             }
         }
 
-        QuickCheck::new().quickcheck(prop as fn(_) -> _)
+        QuickCheck::new().quickcheck(prop as fn(_, _) -> _)
     }
 
     #[test]

From 5e49ff139512485c64e9fcec1f3de34fa324bcac Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:08:19 +1100
Subject: [PATCH 09/20] Don't allocate that many bytes

---
 protocols/gossipsub/src/protocol.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/protocols/gossipsub/src/protocol.rs b/protocols/gossipsub/src/protocol.rs
index 3811055711d..627a42799c7 100644
--- a/protocols/gossipsub/src/protocol.rs
+++ b/protocols/gossipsub/src/protocol.rs
@@ -573,7 +573,7 @@ mod tests {
     #[test]
     /// Test that RPC messages can be encoded and decoded successfully.
     fn encode_decode() {
-        fn prop(message: Message, buf_length: u32) {
+        fn prop(message: Message, buf_length: u16) {
             let message = message.0;
 
             let rpc = Rpc {

From cb90431fac13bfeb7b948a81734eeb1f95261b47 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:08:49 +1100
Subject: [PATCH 10/20] Don't test codec in gossipsub

---
 protocols/gossipsub/src/protocol.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/protocols/gossipsub/src/protocol.rs b/protocols/gossipsub/src/protocol.rs
index 627a42799c7..4cc7253830a 100644
--- a/protocols/gossipsub/src/protocol.rs
+++ b/protocols/gossipsub/src/protocol.rs
@@ -573,7 +573,7 @@ mod tests {
     #[test]
     /// Test that RPC messages can be encoded and decoded successfully.
     fn encode_decode() {
-        fn prop(message: Message, buf_length: u16) {
+        fn prop(message: Message) {
             let message = message.0;
 
             let rpc = Rpc {
@@ -583,7 +583,7 @@ mod tests {
             };
 
             let mut codec = GossipsubCodec::new(u32::MAX as usize, ValidationMode::Strict);
-            let mut buf = BytesMut::with_capacity(buf_length as usize);
+            let mut buf = BytesMut::new();
             codec.encode(rpc.into_protobuf(), &mut buf).unwrap();
             let decoded_rpc = codec.decode(&mut buf).unwrap().unwrap();
             // mark as validated as its a published message

From 51bbb904ac928f18bd410fa04102b1f678976a94 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:19:33 +1100
Subject: [PATCH 11/20] Add test to codec module

---
 Cargo.lock                           |  1 +
 misc/quick-protobuf-codec/Cargo.toml |  1 +
 misc/quick-protobuf-codec/src/lib.rs | 25 +++++++++++++++++++++++++
 protocols/gossipsub/src/protocol.rs  |  2 +-
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index f6efc970bc2..a017111bb29 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4373,6 +4373,7 @@ dependencies = [
  "criterion",
  "futures",
  "quick-protobuf",
+ "quickcheck-ext",
  "thiserror",
  "unsigned-varint 0.8.0",
 ]
diff --git a/misc/quick-protobuf-codec/Cargo.toml b/misc/quick-protobuf-codec/Cargo.toml
index 9c8ed89c231..4ae45900bbc 100644
--- a/misc/quick-protobuf-codec/Cargo.toml
+++ b/misc/quick-protobuf-codec/Cargo.toml
@@ -20,6 +20,7 @@ quick-protobuf = "0.8"
 [dev-dependencies]
 criterion = "0.5.1"
 futures = "0.3.28"
+quickcheck = { workspace = true }
 
 [[bench]]
 name = "codec"
diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 98a16027af6..95f657ef97c 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -180,9 +180,11 @@ impl From<Error> for io::Error {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::proto;
     use asynchronous_codec::FramedRead;
     use futures::io::Cursor;
     use futures::{FutureExt, StreamExt};
+    use quickcheck::{Arbitrary, Gen, QuickCheck};
     use std::error::Error;
 
     #[test]
@@ -199,6 +201,21 @@ mod tests {
         )
     }
 
+    #[test]
+    fn handles_arbitrary_initial_capacity() {
+        fn prop(message: proto::Message, initial_capacity: u16) {
+            let mut buffer = BytesMut::with_capacity(initial_capacity as usize);
+            let mut codec = Codec::<proto::Message>::new(u32::MAX as usize);
+
+            codec.encode(message.clone(), &mut buffer).unwrap();
+            let decoded = codec.decode(&mut buffer).unwrap().unwrap();
+
+            assert_eq!(message, decoded);
+        }
+
+        QuickCheck::new().quickcheck(prop as fn(_, _) -> _)
+    }
+
     /// Constructs a [`BytesMut`] of the provided length where the message is all zeros.
     fn varint_zeroes(length: usize) -> BytesMut {
         let mut buf = unsigned_varint::encode::usize_buffer();
@@ -210,6 +227,14 @@ mod tests {
         src
     }
 
+    impl Arbitrary for proto::Message {
+        fn arbitrary(g: &mut Gen) -> Self {
+            Self {
+                data: Vec::arbitrary(g),
+            }
+        }
+    }
+
     #[derive(Debug)]
     struct Dummy;
 
diff --git a/protocols/gossipsub/src/protocol.rs b/protocols/gossipsub/src/protocol.rs
index 4cc7253830a..42d43c97510 100644
--- a/protocols/gossipsub/src/protocol.rs
+++ b/protocols/gossipsub/src/protocol.rs
@@ -597,7 +597,7 @@ mod tests {
             }
         }
 
-        QuickCheck::new().quickcheck(prop as fn(_, _) -> _)
+        QuickCheck::new().quickcheck(prop as fn(_) -> _)
     }
 
     #[test]

From 30d57ab5f8643c603c335eaec91b39f94f3adce6 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:21:16 +1100
Subject: [PATCH 12/20] Move utilities to the bottom

---
 misc/quick-protobuf-codec/src/lib.rs | 72 ++++++++++++++--------------
 1 file changed, 36 insertions(+), 36 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 95f657ef97c..3da8cb09b7d 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -77,6 +77,42 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
     }
 }
 
+impl<In, Out> Decoder for Codec<In, Out>
+where
+    Out: for<'a> MessageRead<'a>,
+{
+    type Item = Out;
+    type Error = Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        let (len, remaining) = match unsigned_varint::decode::usize(src) {
+            Ok((len, remaining)) => (len, remaining),
+            Err(unsigned_varint::decode::Error::Insufficient) => return Ok(None),
+            Err(e) => return Err(Error(io::Error::new(io::ErrorKind::InvalidData, e))),
+        };
+        let consumed = src.len() - remaining.len();
+        src.advance(consumed);
+
+        if len > self.max_message_len_bytes {
+            return Err(Error(io::Error::new(
+                io::ErrorKind::PermissionDenied,
+                format!(
+                    "message with {len}b exceeds maximum of {}b",
+                    self.max_message_len_bytes
+                ),
+            )));
+        }
+
+        let msg = src.split_to(len);
+
+        let mut reader = BytesReader::from_bytes(&msg);
+        let message = Self::Item::from_reader(&mut reader, &msg)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
+        Ok(Some(message))
+    }
+}
+
 struct UninitMemoryWriterBackend<'a> {
     memory: &'a mut [MaybeUninit<u8>],
     written: &'a mut usize,
@@ -131,42 +167,6 @@ impl<'a> WriterBackend for UninitMemoryWriterBackend<'a> {
     }
 }
 
-impl<In, Out> Decoder for Codec<In, Out>
-where
-    Out: for<'a> MessageRead<'a>,
-{
-    type Item = Out;
-    type Error = Error;
-
-    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
-        let (len, remaining) = match unsigned_varint::decode::usize(src) {
-            Ok((len, remaining)) => (len, remaining),
-            Err(unsigned_varint::decode::Error::Insufficient) => return Ok(None),
-            Err(e) => return Err(Error(io::Error::new(io::ErrorKind::InvalidData, e))),
-        };
-        let consumed = src.len() - remaining.len();
-        src.advance(consumed);
-
-        if len > self.max_message_len_bytes {
-            return Err(Error(io::Error::new(
-                io::ErrorKind::PermissionDenied,
-                format!(
-                    "message with {len}b exceeds maximum of {}b",
-                    self.max_message_len_bytes
-                ),
-            )));
-        }
-
-        let msg = src.split_to(len);
-
-        let mut reader = BytesReader::from_bytes(&msg);
-        let message = Self::Item::from_reader(&mut reader, &msg)
-            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
-
-        Ok(Some(message))
-    }
-}
-
 #[derive(thiserror::Error, Debug)]
 #[error("Failed to encode/decode message")]
 pub struct Error(#[from] io::Error);

From 544130e147cf26aae3fd84240bfb6170f8b1f4fe Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:22:24 +1100
Subject: [PATCH 13/20] Improve naming

---
 misc/quick-protobuf-codec/src/lib.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 3da8cb09b7d..9a0da1016a6 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -52,7 +52,7 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
 
         let mut written = 0;
 
-        let mut writer = Writer::new(UninitMemoryWriterBackend::new(
+        let mut writer = Writer::new(MaybeUninitWriterBackend::new(
             dst.spare_capacity_mut(),
             &mut written,
         ));
@@ -113,18 +113,18 @@ where
     }
 }
 
-struct UninitMemoryWriterBackend<'a> {
+struct MaybeUninitWriterBackend<'a> {
     memory: &'a mut [MaybeUninit<u8>],
     written: &'a mut usize,
 }
 
-impl<'a> UninitMemoryWriterBackend<'a> {
+impl<'a> MaybeUninitWriterBackend<'a> {
     fn new(memory: &'a mut [MaybeUninit<u8>], written: &'a mut usize) -> Self {
         Self { memory, written }
     }
 }
 
-impl<'a> WriterBackend for UninitMemoryWriterBackend<'a> {
+impl<'a> WriterBackend for MaybeUninitWriterBackend<'a> {
     fn pb_write_u8(&mut self, x: u8) -> quick_protobuf::Result<()> {
         self.pb_write_all(&[x])
     }

From fb10b17e5659f19e65ddc4bb70c1d3accfbbed56 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 16:31:04 +1100
Subject: [PATCH 14/20] Split functions to improve readability around unsafe
 code

---
 misc/quick-protobuf-codec/src/lib.rs | 53 ++++++++++++++++++----------
 1 file changed, 34 insertions(+), 19 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 9a0da1016a6..8b845647e11 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -39,42 +39,57 @@ impl<In: MessageWrite, Out> Encoder for Codec<In, Out> {
     type Error = Error;
 
     fn encode(&mut self, item: Self::Item<'_>, dst: &mut BytesMut) -> Result<(), Self::Error> {
-        let message_length = item.get_size();
+        write_length(&item, dst);
+        write_message(&item, dst)?;
 
-        let mut uvi_buf = unsigned_varint::encode::usize_buffer();
-        let encoded_length = unsigned_varint::encode::usize(message_length, &mut uvi_buf);
+        Ok(())
+    }
+}
+
+/// Write the message's length (i.e. `size`) to `dst` as a variable-length integer.
+fn write_length(message: &impl MessageWrite, dst: &mut BytesMut) {
+    let message_length = message.get_size();
+
+    let mut uvi_buf = unsigned_varint::encode::usize_buffer();
+    let encoded_length = unsigned_varint::encode::usize(message_length, &mut uvi_buf);
+
+    dst.extend_from_slice(encoded_length);
+}
 
-        // Append the 'unsigned varint'-encoded length.
-        dst.extend_from_slice(encoded_length);
+/// Write the message itself to `dst`.
+fn write_message(item: &impl MessageWrite, dst: &mut BytesMut) -> io::Result<()> {
+    let message_length = item.get_size();
 
-        // Ensure we have enough capacity to encode our message.
-        dst.reserve(message_length);
+    // Ensure we have enough capacity to encode our message.
+    dst.reserve(message_length);
 
-        let mut written = 0;
+    let mut written = 0;
 
-        let mut writer = Writer::new(MaybeUninitWriterBackend::new(
-            dst.spare_capacity_mut(),
-            &mut written,
-        ));
-        item.write_message(&mut writer)
-            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+    let mut writer = Writer::new(MaybeUninitWriterBackend::new(
+        dst.spare_capacity_mut(),
+        &mut written,
+    ));
+    item.write_message(&mut writer)
+        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
 
+    // Check that we have written exactly as much as we intended to.
+    {
         if written != message_length {
-            return Err(Error(io::Error::new(
+            return Err(io::Error::new(
                 io::ErrorKind::Other,
                 format!(
                     "expected message to be {message_length} bytes long but was {written} bytes"
                 ),
-            )));
+            ));
         }
 
-        // SAFETY: `written` records exactly how many bytes we wrote to `dst`, thus it is safe to extend the length by `written`.
+        // SAFETY: `written` records exactly how many bytes we wrote, hence set them as initialized.
         unsafe {
             dst.set_len(dst.len() + written);
         }
-
-        Ok(())
     }
+
+    Ok(())
 }
 
 impl<In, Out> Decoder for Codec<In, Out>

From 36bc963f00be05cbea128a70f4e2e32d7bcdd350 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 17:05:02 +1100
Subject: [PATCH 15/20] Add more unit tests

---
 misc/quick-protobuf-codec/src/lib.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 8b845647e11..d55b0b7e23b 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -216,6 +216,27 @@ mod tests {
         )
     }
 
+    #[test]
+    fn empty_bytes_mut_does_not_panic() {
+        let mut codec = Codec::<Dummy>::new(100);
+
+        let mut src = varint_zeroes(100);
+        src.truncate(50);
+
+        let result = codec.decode(&mut src);
+
+        assert!(result.unwrap().is_none());
+    }
+
+    #[test]
+    fn only_partial_message_in_bytes_mut_does_not_panic() {
+        let mut codec = Codec::<Dummy>::new(100);
+
+        let result = codec.decode(&mut BytesMut::new());
+
+        assert!(result.unwrap().is_none());
+    }
+
     #[test]
     fn handles_arbitrary_initial_capacity() {
         fn prop(message: proto::Message, initial_capacity: u16) {

From 6cf36420b3e2b9d47aaa8104be5fc18bb9b0b0a8 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 17:08:58 +1100
Subject: [PATCH 16/20] Refactor `decode` for clarity

---
 misc/quick-protobuf-codec/src/lib.rs | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index d55b0b7e23b..1e60c885d66 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -100,15 +100,13 @@ where
     type Error = Error;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
-        let (len, remaining) = match unsigned_varint::decode::usize(src) {
+        let (message_length, remaining) = match unsigned_varint::decode::usize(src) {
             Ok((len, remaining)) => (len, remaining),
             Err(unsigned_varint::decode::Error::Insufficient) => return Ok(None),
             Err(e) => return Err(Error(io::Error::new(io::ErrorKind::InvalidData, e))),
         };
-        let consumed = src.len() - remaining.len();
-        src.advance(consumed);
 
-        if len > self.max_message_len_bytes {
+        if message_length > self.max_message_len_bytes {
             return Err(Error(io::Error::new(
                 io::ErrorKind::PermissionDenied,
                 format!(
@@ -118,10 +116,21 @@ where
             )));
         }
 
-        let msg = src.split_to(len);
+        // Compute how many bytes the varint itself consumed.
+        let varint_length = src.len() - remaining.len();
 
-        let mut reader = BytesReader::from_bytes(&msg);
-        let message = Self::Item::from_reader(&mut reader, &msg)
+        // Ensure we can read an entire message.
+        if src.len() < (message_length + varint_length) {
+            return Ok(None);
+        }
+
+        // Safe to advance buffer now.
+        src.advance(varint_length);
+
+        let message = src.split_to(message_length);
+
+        let mut reader = BytesReader::from_bytes(&message);
+        let message = Self::Item::from_reader(&mut reader, &message)
             .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
 
         Ok(Some(message))
@@ -226,6 +235,11 @@ mod tests {
         let result = codec.decode(&mut src);
 
         assert!(result.unwrap().is_none());
+        assert_eq!(
+            src.len(),
+            50,
+            "to not modify `src` if we cannot read a full message"
+        )
     }
 
     #[test]

From 6eb2e55be05abac44b594d0deb229397392af9f9 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 2 Nov 2023 17:15:18 +1100
Subject: [PATCH 17/20] Fix compile error

---
 misc/quick-protobuf-codec/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 1e60c885d66..61d4d196e3b 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -110,7 +110,7 @@ where
             return Err(Error(io::Error::new(
                 io::ErrorKind::PermissionDenied,
                 format!(
-                    "message with {len}b exceeds maximum of {}b",
+                    "message with {message_length}b exceeds maximum of {}b",
                     self.max_message_len_bytes
                 ),
             )));

From 7893b81bf49c6af297754a3da3aed6b8000859b9 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 9 Nov 2023 13:10:30 +1100
Subject: [PATCH 18/20] Write slice in single command

---
 misc/quick-protobuf-codec/src/lib.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index 61d4d196e3b..da5a8aca1a2 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -182,10 +182,14 @@ impl<'a> WriterBackend for MaybeUninitWriterBackend<'a> {
             return Err(quick_protobuf::errors::Error::UnexpectedEndOfBuffer);
         }
 
-        for b in buf {
-            self.memory[*self.written].write(*b);
-            *self.written += 1;
-        }
+        // SAFETY: &[u8] and &[MaybeUninit<u8>] have the same layout
+        let uninit_src: &[MaybeUninit<u8>] = unsafe { std::mem::transmute(buf) };
+
+        let start = *self.written;
+        let end = *self.written + buf.len();
+
+        self.memory[start..end].copy_from_slice(uninit_src);
+        *self.written += buf.len();
 
         Ok(())
     }

From 9e19b846eac71e4a396622f8460109560ac05e11 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Wed, 22 Nov 2023 10:06:01 +1100
Subject: [PATCH 19/20] Refactor to use `BytesMut` APIs

---
 misc/quick-protobuf-codec/src/lib.rs | 86 ++++++++++------------------
 1 file changed, 30 insertions(+), 56 deletions(-)

diff --git a/misc/quick-protobuf-codec/src/lib.rs b/misc/quick-protobuf-codec/src/lib.rs
index da5a8aca1a2..c50b1264af6 100644
--- a/misc/quick-protobuf-codec/src/lib.rs
+++ b/misc/quick-protobuf-codec/src/lib.rs
@@ -1,11 +1,10 @@
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
 use asynchronous_codec::{Decoder, Encoder};
-use bytes::{Buf, BytesMut};
+use bytes::{Buf, BufMut, BytesMut};
 use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer, WriterBackend};
 use std::io;
 use std::marker::PhantomData;
-use std::mem::MaybeUninit;
 
 mod generated;
 
@@ -58,37 +57,10 @@ fn write_length(message: &impl MessageWrite, dst: &mut BytesMut) {
 
 /// Write the message itself to `dst`.
 fn write_message(item: &impl MessageWrite, dst: &mut BytesMut) -> io::Result<()> {
-    let message_length = item.get_size();
-
-    // Ensure we have enough capacity to encode our message.
-    dst.reserve(message_length);
-
-    let mut written = 0;
-
-    let mut writer = Writer::new(MaybeUninitWriterBackend::new(
-        dst.spare_capacity_mut(),
-        &mut written,
-    ));
+    let mut writer = Writer::new(BytesMutWriterBackend::new(dst));
     item.write_message(&mut writer)
         .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
 
-    // Check that we have written exactly as much as we intended to.
-    {
-        if written != message_length {
-            return Err(io::Error::new(
-                io::ErrorKind::Other,
-                format!(
-                    "expected message to be {message_length} bytes long but was {written} bytes"
-                ),
-            ));
-        }
-
-        // SAFETY: `written` records exactly how many bytes we wrote, hence set them as initialized.
-        unsafe {
-            dst.set_len(dst.len() + written);
-        }
-    }
-
     Ok(())
 }
 
@@ -137,59 +109,61 @@ where
     }
 }
 
-struct MaybeUninitWriterBackend<'a> {
-    memory: &'a mut [MaybeUninit<u8>],
-    written: &'a mut usize,
+struct BytesMutWriterBackend<'a> {
+    dst: &'a mut BytesMut,
 }
 
-impl<'a> MaybeUninitWriterBackend<'a> {
-    fn new(memory: &'a mut [MaybeUninit<u8>], written: &'a mut usize) -> Self {
-        Self { memory, written }
+impl<'a> BytesMutWriterBackend<'a> {
+    fn new(dst: &'a mut BytesMut) -> Self {
+        Self { dst }
     }
 }
 
-impl<'a> WriterBackend for MaybeUninitWriterBackend<'a> {
+impl<'a> WriterBackend for BytesMutWriterBackend<'a> {
     fn pb_write_u8(&mut self, x: u8) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&[x])
+        self.dst.put_u8(x);
+
+        Ok(())
     }
 
     fn pb_write_u32(&mut self, x: u32) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_u32_le(x);
+
+        Ok(())
     }
 
     fn pb_write_i32(&mut self, x: i32) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_i32_le(x);
+
+        Ok(())
     }
 
     fn pb_write_f32(&mut self, x: f32) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_f32_le(x);
+
+        Ok(())
     }
 
     fn pb_write_u64(&mut self, x: u64) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_u64_le(x);
+
+        Ok(())
     }
 
     fn pb_write_i64(&mut self, x: i64) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_i64_le(x);
+
+        Ok(())
     }
 
     fn pb_write_f64(&mut self, x: f64) -> quick_protobuf::Result<()> {
-        self.pb_write_all(&x.to_le_bytes())
+        self.dst.put_f64_le(x);
+
+        Ok(())
     }
 
     fn pb_write_all(&mut self, buf: &[u8]) -> quick_protobuf::Result<()> {
-        if self.memory.len() - *self.written < buf.len() {
-            return Err(quick_protobuf::errors::Error::UnexpectedEndOfBuffer);
-        }
-
-        // SAFETY: &[u8] and &[MaybeUninit<u8>] have the same layout
-        let uninit_src: &[MaybeUninit<u8>] = unsafe { std::mem::transmute(buf) };
-
-        let start = *self.written;
-        let end = *self.written + buf.len();
-
-        self.memory[start..end].copy_from_slice(uninit_src);
-        *self.written += buf.len();
+        self.dst.put_slice(buf);
 
         Ok(())
     }

From 5514f2ed24d727290d8628cb20c58be427bcdaeb Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Wed, 22 Nov 2023 10:14:42 +1100
Subject: [PATCH 20/20] Update misc/quick-protobuf-codec/CHANGELOG.md

---
 misc/quick-protobuf-codec/CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/misc/quick-protobuf-codec/CHANGELOG.md b/misc/quick-protobuf-codec/CHANGELOG.md
index bd81c3d589b..a301293621f 100644
--- a/misc/quick-protobuf-codec/CHANGELOG.md
+++ b/misc/quick-protobuf-codec/CHANGELOG.md
@@ -1,4 +1,4 @@
-## 0.3.1 - unreleased
+## 0.3.1
 
 - Reduce allocations during encoding.
   See [PR 4782](https://github.com/libp2p/rust-libp2p/pull/4782).