Skip to content

Commit

Permalink
Handle unaligned output instead of using UB unsafe allocs
Browse files Browse the repository at this point in the history
- Vec::from_raw_parts is much stricter than anticipated and requires
  deallocation to happen with the same alignment
  • Loading branch information
lwus committed May 11, 2022
1 parent 58e8a6c commit 1382308
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 32 deletions.
26 changes: 10 additions & 16 deletions src/decode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,20 +227,9 @@ impl<'a, I: AsRef<[u8]>> DecodeBuilder<'a, I> {
#[cfg(feature = "alloc")]
#[cfg_attr(docsrs, doc(cfg(any(feature = "alloc", feature = "std"))))]
pub fn into_vec_unsafe(self) -> Result<Vec<u8>> {
let mut output = Vec::<u32>::new();
let mut output = Vec::new();
output.resize((self.input.as_ref().len() + 3) / 4 * 4, 0);

// Prevent running `output`'s destructor so we are in complete control
// of the allocation.
let mut output = std::mem::ManuallyDrop::new(output);

// Pull out the various important pieces of information about `output`
let p = output.as_mut_ptr();
let len = output.len();
let cap = output.capacity();

let mut output = unsafe { Vec::<u8>::from_raw_parts(p as *mut u8, len * 4, cap * 4) };

let len = decode_into_limbs(self.input.as_ref(), &mut output, self.alpha)?;
output.truncate(len);
Ok(output)
Expand Down Expand Up @@ -356,10 +345,7 @@ fn decode_into_limbs(input: &[u8], output: &mut [u8], alpha: &Alphabet) -> Resul
let next_limb_multiplier = 58 * 58 * 58 * 58 * 58;

let (prefix, output_as_limbs, _) = unsafe { output.align_to_mut::<u32>() };
if prefix.len() != 0 {
// invariant
return Err(Error::BufferTooSmall);
}
let prefix_len = prefix.len();

while input_iter.len() >= input_bytes_per_limb {
let input_byte0 = decode_input_byte(input_iter.next().unwrap())?;
Expand Down Expand Up @@ -414,6 +400,8 @@ fn decode_into_limbs(input: &[u8], output: &mut [u8], alpha: &Alphabet) -> Resul

// rescale for the remainder
index = index * 4;
{
let output = &mut output[prefix_len..];
while index > 0 && output[index - 1] == 0 {
index -= 1;
}
Expand All @@ -426,6 +414,12 @@ fn decode_into_limbs(input: &[u8], output: &mut [u8], alpha: &Alphabet) -> Resul
}

output[..index].reverse();
}

if prefix_len > 0 {
output.copy_within(prefix_len..prefix_len + index, 0);
}

Ok(index)
}

Expand Down
26 changes: 10 additions & 16 deletions src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -250,21 +250,10 @@ impl<'a, I: AsRef<[u8]>> EncodeBuilder<'a, I> {

/// Encode into a new owned vector.
pub fn into_vec_unsafe(self) -> Vec<u8> {
let mut output = Vec::<u32>::new();
let mut output = Vec::new();
let max_encoded_len = (self.input.as_ref().len() / 5 + 1) * 8;
output.resize((max_encoded_len + 3) / 4 * 4, 0);

// Prevent running `output`'s destructor so we are in complete control
// of the allocation.
let mut output = std::mem::ManuallyDrop::new(output);

// Pull out the various important pieces of information about `output`
let p = output.as_mut_ptr();
let len = output.len();
let cap = output.capacity();

let mut output = unsafe { Vec::<u8>::from_raw_parts(p as *mut u8, len * 4, cap * 4) };

let len = encode_into_limbs(self.input.as_ref(), &mut output, self.alpha).unwrap();
output.truncate(len);
output
Expand Down Expand Up @@ -399,10 +388,7 @@ where
{
let input_bytes_per_limb = 3;
let (prefix, output_as_limbs, _) = unsafe { output.align_to_mut::<u32>() };
if prefix.len() != 0 {
// invariant
return Err(Error::BufferTooSmall);
}
let prefix_len = prefix.len();

let mut index = 0;
let mut input_iter = input.clone().into_iter();
Expand Down Expand Up @@ -470,6 +456,8 @@ where

// rescale for the remainder
index = index * 4;
{
let output = &mut output[prefix_len..];
while index > 0 && output[index - 1] == 0 {
index -= 1;
}
Expand All @@ -485,6 +473,12 @@ where
}

output[..index].reverse();
}

if prefix_len > 0 {
output.copy_within(prefix_len..prefix_len + index, 0);
}

Ok(index)
}

Expand Down

0 comments on commit 1382308

Please sign in to comment.