Introduce API to safely initialize Packets

ci/test-miri.sh

@@ -6,8 +6,8 @@ source ci/_
 source ci/rust-version.sh nightly
 
 # miri is very slow; so only run very few of selective tests!
+_ cargo "+${rust_nightly}" miri test -p solana-packet -- test_packet_buffer_writer
 _ cargo "+${rust_nightly}" miri test -p solana-program -- hash:: account_info::
-
 _ cargo "+${rust_nightly}" miri test -p solana-unified-scheduler-logic
 
 # run intentionally-#[ignored] ub triggering tests for each to make sure they fail

core/src/shred_fetch_stage.rs

@@ -14,7 +14,7 @@ use {
         clock::{Slot, DEFAULT_MS_PER_SLOT},
         epoch_schedule::EpochSchedule,
         genesis_config::ClusterType,
-        packet::{Meta, PACKET_DATA_SIZE},
+        packet::{Packet, PACKET_DATA_SIZE},
         pubkey::Pubkey,
     },
     solana_streamer::streamer::{self, PacketBatchReceiver, StreamerReceiveStats},
@@ -350,29 +350,23 @@ pub(crate) fn receive_quic_datagrams(
         };
         let mut packet_batch =
             PacketBatch::new_with_recycler(&recycler, PACKETS_PER_BATCH, "receive_quic_datagrams");
-        unsafe {
-            packet_batch.set_len(PACKETS_PER_BATCH);
-        };
         let deadline = Instant::now() + PACKET_COALESCE_DURATION;
         let entries = std::iter::once(entry).chain(
             std::iter::repeat_with(|| quic_datagrams_receiver.recv_deadline(deadline).ok())
                 .while_some(),
         );
         let size = entries
             .filter(|(_, _, bytes)| bytes.len() <= PACKET_DATA_SIZE)
-            .zip(packet_batch.iter_mut())
+            .zip(packet_batch.spare_capacity_mut().iter_mut())
             .map(|((_pubkey, addr, bytes), packet)| {
-                *packet.meta_mut() = Meta {
-                    size: bytes.len(),
-                    addr: addr.ip(),
-                    port: addr.port(),
-                    flags,
-                };
-                packet.buffer_mut()[..bytes.len()].copy_from_slice(&bytes);
+                Packet::init_packet_from_bytes(packet, &bytes, Some(&addr)).unwrap();
+                // SAFETY: Packet::init_packet_from_bytes() just initialized the packet
+                unsafe { packet.assume_init_mut().meta_mut().set_flags(flags) };
             })
             .count();
         if size > 0 {
-            packet_batch.truncate(size);
+            // SAFETY: By now, size packets have been initialized
+            unsafe { packet_batch.set_len(size) };
             if sender.send(packet_batch).is_err() {
                 return; // The receiver end of the channel is disconnected.
             }

entry/src/entry.rs

@@ -24,7 +24,6 @@ use {
     solana_runtime_transaction::runtime_transaction::RuntimeTransaction,
     solana_sdk::{
         hash::Hash,
-        packet::Meta,
         transaction::{
             Result, SanitizedTransaction, Transaction, TransactionError,
             TransactionVerificationMode, VersionedTransaction,
@@ -548,26 +547,22 @@ fn start_verify_transactions_gpu(
                     num_transactions,
                     "entry-sig-verify",
                 );
-                // We use set_len here instead of resize(num_txs, Packet::default()), to save
-                // memory bandwidth and avoid writing a large amount of data that will be overwritten
-                // soon afterwards. As well, Packet::default() actually leaves the packet data
-                // uninitialized, so the initialization would simply write junk into
-                // the vector anyway.
-                unsafe {
-                    packet_batch.set_len(num_transactions);
-                }
+
+                let uninitialized_packets = packet_batch.spare_capacity_mut().iter_mut();
                 let transaction_iter = transaction_chunk
                     .iter()
                     .map(|tx| tx.to_versioned_transaction());
 
-                let res = packet_batch
-                    .iter_mut()
-                    .zip(transaction_iter)
-                    .all(|(packet, tx)| {
-                        *packet.meta_mut() = Meta::default();
-                        Packet::populate_packet(packet, None, &tx).is_ok()
-                    });
-                if res {
+                let all_packets_initialized =
+                    uninitialized_packets
+                        .zip(transaction_iter)
+                        .all(|(uninit_packet, tx)| {
+                            Packet::init_packet_from_data(uninit_packet, &tx, None).is_ok()
+                        });
+
+                if all_packets_initialized {
+                    // SAFETY: All packets have been successfully initialized
+                    unsafe { packet_batch.set_len(num_transactions) };
                     Ok(packet_batch)
                 } else {
                     Err(TransactionError::SanitizeFailure)

perf/src/cuda_runtime.rs

@@ -14,6 +14,7 @@ use {
     rayon::prelude::*,
     serde::{Deserialize, Serialize},
     std::{
+        mem::MaybeUninit,
         ops::{Index, IndexMut},
         os::raw::c_int,
         slice::{Iter, IterMut, SliceIndex},
@@ -129,6 +130,10 @@ impl<T: Clone + Default + Sized> PinnedVec<T> {
         self.x.iter_mut()
     }
 
+    pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
+        self.x.spare_capacity_mut()
+    }
+
     pub fn capacity(&self) -> usize {
         self.x.capacity()
     }

perf/src/packet.rs

@@ -7,6 +7,7 @@ use {
     serde::{de::DeserializeOwned, Deserialize, Serialize},
     std::{
         io::Read,
+        mem::MaybeUninit,
         net::SocketAddr,
         ops::{Index, IndexMut},
         slice::{Iter, IterMut, SliceIndex},
@@ -152,6 +153,10 @@ impl PacketBatch {
         self.packets.iter_mut()
     }
 
+    pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<Packet>] {
+        self.packets.spare_capacity_mut()
+    }
+
     /// See Vector::set_len() for more details
     ///
     /// # Safety

sdk/packet/src/lib.rs

@@ -10,7 +10,10 @@ use {
     bitflags::bitflags,
     std::{
         fmt,
+        io::{self, Write},
+        mem,
         net::{IpAddr, Ipv4Addr, SocketAddr},
+        ptr,
         slice::SliceIndex,
     },
 };
@@ -158,11 +161,76 @@ impl Packet {
         &mut self.meta
     }
 
+    #[cfg(feature = "bincode")]
+    /// Initializes a std::mem::MaybeUninit<Packet> such that the Packet can
+    /// be safely extracted via methods such as MaybeUninit::assume_init()
+    pub fn init_packet_from_data<T: serde::Serialize>(
+        packet: &mut mem::MaybeUninit<Packet>,
+        data: &T,
+        addr: Option<&SocketAddr>,
+    ) -> Result<()> {
+        let mut writer = PacketWriter::new_from_uninit_packet(packet);
+        bincode::serialize_into(&mut writer, data)?;
+
+        let serialized_size = writer.position();
+        let (ip, port) = if let Some(addr) = addr {
+            (addr.ip(), addr.port())
+        } else {
+            (IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
+        };
+        Self::init_packet_meta(
+            packet,
+            Meta {
+                size: serialized_size,
+                addr: ip,
+                port,
+                flags: PacketFlags::empty(),
+            },
+        );
+
+        Ok(())
+    }
+
+    pub fn init_packet_from_bytes(
+        packet: &mut mem::MaybeUninit<Packet>,
+        bytes: &[u8],
+        addr: Option<&SocketAddr>,
+    ) -> io::Result<()> {
+        let mut writer = PacketWriter::new_from_uninit_packet(packet);
+        let num_bytes_written = writer.write(bytes)?;
+        debug_assert_eq!(bytes.len(), num_bytes_written);
+
+        let size = writer.position();
+        let (ip, port) = if let Some(addr) = addr {
+            (addr.ip(), addr.port())
+        } else {
+            (IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
+        };
+        Self::init_packet_meta(
+            packet,
+            Meta {
+                size,
+                addr: ip,
+                port,
+                flags: PacketFlags::empty(),
+            },
+        );
+
+        Ok(())
+    }
+
+    fn init_packet_meta(packet: &mut mem::MaybeUninit<Packet>, meta: Meta) {
+        // SAFETY: Access the field by pointer as creating a reference to
+        // and/or within the uninitialized Packet is undefined behavior
+        unsafe { ptr::addr_of_mut!((*packet.as_mut_ptr()).meta).write(meta) };
+    }
+
     #[cfg(feature = "bincode")]
     pub fn from_data<T: serde::Serialize>(dest: Option<&SocketAddr>, data: T) -> Result<Self> {
-        let mut packet = Self::default();
-        Self::populate_packet(&mut packet, dest, &data)?;
-        Ok(packet)
+        let mut packet = mem::MaybeUninit::uninit();
+        Self::init_packet_from_data(&mut packet, &data, dest)?;
+        // SAFETY: init_packet_from_data() just initialized the packet
+        unsafe { Ok(packet.assume_init()) }
     }
 
     #[cfg(feature = "bincode")]
@@ -224,6 +292,61 @@ impl PartialEq for Packet {
     }
 }
 
+/// A custom implementation of io::Write to facilitate safe (non-UB)
+/// initialization of a MaybeUninit<Packet>
+struct PacketWriter {
+    // A pointer to the current write position
+    position: *mut u8,
+    // The number of remaining bytes that can be written to
+    spare_capacity: usize,
+}
+
+impl PacketWriter {
+    fn new_from_uninit_packet(packet: &mut mem::MaybeUninit<Packet>) -> Self {
+        // SAFETY: Access the field by pointer as creating a reference to
+        // and/or within the uninitialized Packet is undefined behavior
+        let position = unsafe { ptr::addr_of_mut!((*packet.as_mut_ptr()).buffer) as *mut u8 };
+        let spare_capacity = PACKET_DATA_SIZE;
+
+        Self {
+            position,
+            spare_capacity,
+        }
+    }
+
+    /// The offset of the write pointer within the buffer, which is also the
+    /// number of bytes that have been written
+    fn position(&self) -> usize {
+        PACKET_DATA_SIZE.saturating_sub(self.spare_capacity)
+    }
+}
+
+impl io::Write for PacketWriter {
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if buf.len() > self.spare_capacity {
+            return Err(io::Error::from(io::ErrorKind::WriteZero));
+        }
+
+        // SAFETY: We previously verifed that buf.len() <= self.spare_capacity
+        // so this write will not push us past the end of the buffer. Likewise,
+        // we can update self.spare_capacity without fear of overflow
+        unsafe {
+            ptr::copy(buf.as_ptr(), self.position, buf.len());
+            // Update position and spare_capacity for the next call to write()
+            self.position = self.position.add(buf.len());
+            self.spare_capacity = self.spare_capacity.saturating_sub(buf.len());
+        }
+
+        Ok(buf.len())
+    }
+
+    #[inline]
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
 impl Meta {
     pub fn socket_addr(&self) -> SocketAddr {
         SocketAddr::new(self.addr, self.port)
@@ -239,6 +362,11 @@ impl Meta {
             .set(PacketFlags::FROM_STAKED_NODE, from_staked_node);
     }
 
+    #[inline]
+    pub fn set_flags(&mut self, flags: PacketFlags) {
+        self.flags = flags;
+    }
+
     #[inline]
     pub fn discard(&self) -> bool {
         self.flags.contains(PacketFlags::DISCARD)
@@ -309,7 +437,7 @@ impl Default for Meta {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use {super::*, std::io::Write};
 
     #[test]
     fn test_deserialize_slice() {
@@ -349,4 +477,35 @@ mod tests {
             Err("the size limit has been reached".to_string()),
         );
     }
+
+    #[test]
+    fn test_packet_buffer_writer() {
+        let mut packet = mem::MaybeUninit::<Packet>::uninit();
+        let mut writer = PacketWriter::new_from_uninit_packet(&mut packet);
+        let total_capacity = writer.spare_capacity;
+        assert_eq!(total_capacity, PACKET_DATA_SIZE);
+        let payload: [u8; PACKET_DATA_SIZE] = std::array::from_fn(|i| i as u8 % 255);
+
+        // Write 1200 bytes (1200 total)
+        let num_to_write = PACKET_DATA_SIZE - 32;
+        assert_eq!(
+            num_to_write,
+            writer.write(&payload[..num_to_write]).unwrap()
+        );
+        assert_eq!(num_to_write, writer.position());
+        // Write 28 bytes (1228 total)
+        assert_eq!(28, writer.write(&payload[1200..1200 + 28]).unwrap());
+        assert_eq!(1200 + 28, writer.position());
+        // Attempt to write 5 bytes (1233 total) which exceeds buffer capacity
+        assert!(writer
+            .write(&payload[1200 + 28 - 1..PACKET_DATA_SIZE])
+            .is_err());
+        // writer.position() remains unchanged
+        assert_eq!(1200 + 28, writer.position());
+        // Write 4 bytes (1232 total) to fill buffer
+        assert_eq!(4, writer.write(&payload[1200 + 28..]).unwrap());
+        assert_eq!(PACKET_DATA_SIZE, writer.position());
+        // Writing any amount of bytes will fail on the already full buffer
+        assert!(writer.write(&[0]).is_err());
+    }
 }