Skip to content

Commit 04905e8

Browse files
404Setup404Setup
authored
Fix buffer length manipulation without initialization (#426)
Replaced unsafe `set_len` calls on uninitialized memory with idiomatic safe patterns using `spare_capacity_mut` and deferred `set_len` after successful initialization. Added safety assertions to prevent potential buffer overflows.
1 parent bc23fbc commit 04905e8

3 files changed

Lines changed: 45 additions & 41 deletions

File tree

‎src/batch.rs‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ impl BatchCompressor {
2626
let (res, size, _) =
2727
compressor.compress(input, buf_slice, crate::compress::FlushMode::Finish);
2828
if res == CompressResult::Success {
29+
assert!(size <= bound);
2930
unsafe {
3031
buffer.set_len(size);
3132
}
@@ -59,16 +60,20 @@ impl BatchDecompressor {
5960
.map_init(
6061
|| (Decompressor::new(), Vec::new()),
6162
|(decompressor, buffer), (&input, &max_size)| {
63+
buffer.clear();
6264
if buffer.capacity() < max_size {
63-
buffer.reserve(max_size.saturating_sub(buffer.len()));
64-
}
65-
unsafe {
66-
buffer.set_len(max_size);
65+
buffer.reserve(max_size);
6766
}
67+
let buf_uninit = buffer.spare_capacity_mut();
68+
let buf_slice = &mut buf_uninit[..max_size];
6869

69-
let (res, _, size) = decompressor.decompress(input, buffer);
70+
let (res, _, size) = unsafe { decompressor.decompress_uninit(input, buf_slice) };
7071
if res == DecompressResult::Success {
71-
Some(buffer[..size].to_vec())
72+
assert!(size <= max_size);
73+
unsafe {
74+
buffer.set_len(size);
75+
}
76+
Some(buffer.to_vec())
7277
} else {
7378
None
7479
}

‎src/compress/mod.rs‎

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -684,17 +684,17 @@ impl Compressor {
684684
let mode = if is_last { flush_mode } else { FlushMode::Sync };
685685

686686
let bound = Self::deflate_compress_bound(chunk.len());
687+
buf.clear();
687688
if buf.capacity() < bound {
688-
buf.reserve(bound - buf.len());
689-
}
690-
unsafe {
691-
buf.set_len(bound);
689+
buf.reserve(bound);
692690
}
693691

694-
let buf_uninit = slice_as_uninit_mut(buf);
692+
let buf_uninit = buf.spare_capacity_mut();
693+
let buf_uninit = &mut buf_uninit[..bound];
695694

696695
let (res, size, _) = compressor.compress(chunk, buf_uninit, mode);
697696
if res == CompressResult::Success {
697+
assert!(size <= bound);
698698
unsafe {
699699
buf.set_len(size);
700700
}
@@ -924,12 +924,7 @@ impl Compressor {
924924
self.dp_costs[0] = 0;
925925

926926
self.dp_path.clear();
927-
if self.dp_path.capacity() < processed + 1 {
928-
self.dp_path.reserve(processed + 1 - self.dp_path.len());
929-
}
930-
unsafe {
931-
self.dp_path.set_len(processed + 1);
932-
}
927+
self.dp_path.resize(processed + 1, 0);
933928

934929
mf.reset();
935930
let mut pos = 0;
@@ -1724,12 +1719,7 @@ impl Compressor {
17241719
self.dp_costs[0] = 0;
17251720

17261721
self.dp_path.clear();
1727-
if self.dp_path.capacity() < processed + 1 {
1728-
self.dp_path.reserve(processed + 1 - self.dp_path.len());
1729-
}
1730-
unsafe {
1731-
self.dp_path.set_len(processed + 1);
1732-
}
1722+
self.dp_path.resize(processed + 1, 0);
17331723

17341724
mf.reset();
17351725
let mut pos = 0;

‎src/stream.rs‎

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -74,23 +74,26 @@ impl<W: Write + Send> DeflateEncoder<W> {
7474
if !final_block {
7575
bound += 5;
7676
}
77-
if output.len() < bound {
77+
output.clear();
78+
if output.capacity() < bound {
7879
output
79-
.try_reserve(bound - output.len())
80+
.try_reserve(bound)
8081
.map_err(io::Error::other)?;
81-
unsafe {
82-
output.set_len(bound);
83-
}
8482
}
8583

8684
let mode = if final_block {
8785
crate::compress::FlushMode::Finish
8886
} else {
8987
crate::compress::FlushMode::Sync
9088
};
91-
let out_uninit = crate::common::slice_as_uninit_mut(output);
89+
let out_uninit = output.spare_capacity_mut();
90+
let out_uninit = &mut out_uninit[..bound];
9291
let (res, size, _) = compressor.compress(chunk, out_uninit, mode);
9392
if res == CompressResult::Success {
93+
assert!(size <= bound);
94+
unsafe {
95+
output.set_len(size);
96+
}
9497
if let Some(writer) = &mut self.writer {
9598
writer.write_all(&output[..size])?;
9699
}
@@ -108,23 +111,26 @@ impl<W: Write + Send> DeflateEncoder<W> {
108111
if !(final_block && i == num_chunks - 1) {
109112
bound += 5;
110113
}
111-
if output.len() < bound {
114+
output.clear();
115+
if output.capacity() < bound {
112116
output
113-
.try_reserve(bound - output.len())
117+
.try_reserve(bound)
114118
.map_err(io::Error::other)?;
115-
unsafe {
116-
output.set_len(bound);
117-
}
118119
}
119120

120121
let mode = if final_block && i == num_chunks - 1 {
121122
crate::compress::FlushMode::Finish
122123
} else {
123124
crate::compress::FlushMode::Sync
124125
};
125-
let out_uninit = crate::common::slice_as_uninit_mut(output);
126+
let out_uninit = output.spare_capacity_mut();
127+
let out_uninit = &mut out_uninit[..bound];
126128
let (res, size, _) = compressor.compress(chunk, out_uninit, mode);
127129
if res == CompressResult::Success {
130+
assert!(size <= bound);
131+
unsafe {
132+
output.set_len(size);
133+
}
128134
Ok(size)
129135
} else {
130136
Err(io::Error::other("Compression failed"))
@@ -154,23 +160,26 @@ impl<W: Write + Send> DeflateEncoder<W> {
154160
if !final_block {
155161
bound += 5;
156162
}
157-
if output.len() < bound {
163+
output.clear();
164+
if output.capacity() < bound {
158165
output
159-
.try_reserve(bound - output.len())
166+
.try_reserve(bound)
160167
.map_err(io::Error::other)?;
161-
unsafe {
162-
output.set_len(bound);
163-
}
164168
}
165169

166170
let mode = if final_block {
167171
crate::compress::FlushMode::Finish
168172
} else {
169173
crate::compress::FlushMode::Sync
170174
};
171-
let out_uninit = crate::common::slice_as_uninit_mut(output);
175+
let out_uninit = output.spare_capacity_mut();
176+
let out_uninit = &mut out_uninit[..bound];
172177
let (res, size, _) = compressor.compress(&self.buffer, out_uninit, mode);
173178
if res == CompressResult::Success {
179+
assert!(size <= bound);
180+
unsafe {
181+
output.set_len(size);
182+
}
174183
if let Some(writer) = &mut self.writer {
175184
writer.write_all(&output[..size])?;
176185
}

0 commit comments

Comments
 (0)