Skip to content

Commit b152ad4

Browse files
TestingPlantandrewgazelka
authored andcommitted
fix: initialize scratch buffer
Previously, the code took a reference to uninitialized memory, which is undefined behavior. Allocating zeroed memory is relatively cheap, so the scratch buffer allocates zeroed memory to initialize the memory.
1 parent 82f32bf commit b152ad4

File tree

2 files changed

+23
-36
lines changed

2 files changed

+23
-36
lines changed

crates/hyperion/src/lib.rs

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -383,26 +383,21 @@ impl HyperionCore {
383383
/// A scratch buffer for intermediate operations. This will return an empty [`Vec`] when calling [`Scratch::obtain`].
384384
#[derive(Debug)]
385385
pub struct Scratch<A: Allocator = std::alloc::Global> {
386-
inner: Vec<u8, A>,
386+
inner: Box<[u8], A>,
387387
}
388388

389389
impl Default for Scratch<std::alloc::Global> {
390390
fn default() -> Self {
391-
let inner = Vec::with_capacity(MAX_PACKET_SIZE);
392-
Self { inner }
391+
std::alloc::Global.into()
393392
}
394393
}
395394

396395
/// Nice for getting a buffer that can be used for intermediate work
397-
///
398-
/// # Safety
399-
/// - every single time [`ScratchBuffer::obtain`] is called, the buffer will be cleared before returning
400-
/// - the buffer has capacity of at least `MAX_PACKET_SIZE`
401-
pub unsafe trait ScratchBuffer: sealed::Sealed + Debug {
396+
pub trait ScratchBuffer: sealed::Sealed + Debug {
402397
/// The type of the allocator the [`Vec`] uses.
403398
type Allocator: Allocator;
404-
/// Obtains a buffer that can be used for intermediate work.
405-
fn obtain(&mut self) -> &mut Vec<u8, Self::Allocator>;
399+
/// Obtains a buffer that can be used for intermediate work. The contents are unspecified.
400+
fn obtain(&mut self) -> &mut [u8];
406401
}
407402

408403
mod sealed {
@@ -411,20 +406,23 @@ mod sealed {
411406

412407
impl<A: Allocator + Debug> sealed::Sealed for Scratch<A> {}
413408

414-
unsafe impl<A: Allocator + Debug> ScratchBuffer for Scratch<A> {
409+
impl<A: Allocator + Debug> ScratchBuffer for Scratch<A> {
415410
type Allocator = A;
416411

417-
fn obtain(&mut self) -> &mut Vec<u8, Self::Allocator> {
418-
self.inner.clear();
412+
fn obtain(&mut self) -> &mut [u8] {
419413
&mut self.inner
420414
}
421415
}
422416

423417
impl<A: Allocator> From<A> for Scratch<A> {
424418
fn from(allocator: A) -> Self {
425-
Self {
426-
inner: Vec::with_capacity_in(MAX_PACKET_SIZE, allocator),
427-
}
419+
// A zeroed slice is allocated to avoid reading from uninitialized memory, which is UB.
420+
// Allocating zeroed memory is usually very cheap, so there are minimal performance
421+
// penalties from this.
422+
let inner = Box::new_zeroed_slice_in(MAX_PACKET_SIZE, allocator);
423+
// SAFETY: The box was initialized to zero, and u8 can be represented by zero
424+
let inner = unsafe { inner.assume_init() };
425+
Self { inner }
428426
}
429427
}
430428

crates/hyperion/src/net/encoder/mod.rs

Lines changed: 9 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
use std::{
1010
fmt::Debug,
1111
io::{Cursor, Write},
12-
mem::MaybeUninit,
1312
};
1413

1514
use anyhow::ensure;
@@ -129,37 +128,27 @@ impl PacketEncoder {
129128
if data_len > threshold {
130129
let scratch = scratch.obtain();
131130

132-
debug_assert!(scratch.is_empty());
133-
134131
let data_slice = &mut slice
135132
[usize::try_from(data_write_start)?..usize::try_from(end_data_position_exclusive)?];
136133

137-
{
138-
// todo: I think this kinda safe maybe??? ... lol. well I know at least scratch is always large enough
139-
let written = {
140-
let scratch = scratch.spare_capacity_mut();
141-
let scratch = unsafe { MaybeUninit::slice_assume_init_mut(scratch) };
142-
143-
let len = data_slice.len();
144-
let span = tracing::trace_span!("zlib_compress", bytes = len);
145-
let _enter = span.enter();
146-
compressor.zlib_compress(data_slice, scratch)?
147-
};
134+
let written = {
135+
let len = data_slice.len();
136+
let span = tracing::trace_span!("zlib_compress", bytes = len);
137+
let _enter = span.enter();
138+
compressor.zlib_compress(data_slice, scratch).unwrap()
139+
};
148140

149-
unsafe {
150-
scratch.set_len(scratch.len() + written);
151-
}
152-
}
141+
let compressed = &scratch[..written];
153142

154143
let data_len = VarInt(data_len as u32 as i32);
155144

156-
let packet_len = data_len.written_size() + scratch.len();
145+
let packet_len = data_len.written_size() + compressed.len();
157146
let packet_len = VarInt(packet_len as u32 as i32);
158147

159148
let mut write = Cursor::new(&mut slice[..]);
160149
packet_len.encode(&mut write)?;
161150
data_len.encode(&mut write)?;
162-
write.write_all(scratch)?;
151+
write.write_all(compressed)?;
163152

164153
let len = write.position();
165154

0 commit comments

Comments
 (0)