Merge pull request #126 from Dushistov/fix-deprecation

Make tungstenite-rs more safe
pull/130/head
Alexey Galakhov 5 years ago committed by GitHub
commit bbaa48b4d6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 98
      src/protocol/frame/mask.rs

@ -1,9 +1,3 @@
use rand;
use std::cmp::min;
#[allow(deprecated)]
use std::mem::uninitialized;
use std::ptr::{copy_nonoverlapping, read_unaligned};
/// Generate a random frame mask. /// Generate a random frame mask.
#[inline] #[inline]
pub fn generate_mask() -> [u8; 4] { pub fn generate_mask() -> [u8; 4] {
@ -18,7 +12,6 @@ pub fn apply_mask(buf: &mut [u8], mask: [u8; 4]) {
/// A safe unoptimized mask application. /// A safe unoptimized mask application.
#[inline] #[inline]
#[allow(dead_code)]
fn apply_mask_fallback(buf: &mut [u8], mask: [u8; 4]) { fn apply_mask_fallback(buf: &mut [u8], mask: [u8; 4]) {
for (i, byte) in buf.iter_mut().enumerate() { for (i, byte) in buf.iter_mut().enumerate() {
*byte ^= mask[i & 3]; *byte ^= mask[i & 3];
@ -27,21 +20,13 @@ fn apply_mask_fallback(buf: &mut [u8], mask: [u8; 4]) {
/// Faster version of `apply_mask()` which operates on 4-byte blocks. /// Faster version of `apply_mask()` which operates on 4-byte blocks.
#[inline] #[inline]
#[allow(dead_code, clippy::cast_ptr_alignment)] pub fn apply_mask_fast32(buf: &mut [u8], mask: [u8; 4]) {
fn apply_mask_fast32(buf: &mut [u8], mask: [u8; 4]) { let mask_u32 = u32::from_ne_bytes(mask);
let mask_u32: u32 = unsafe { read_unaligned(mask.as_ptr() as *const u32) };
let mut ptr = buf.as_mut_ptr();
let mut len = buf.len();
// Possible first unaligned block. let (mut prefix, words, mut suffix) = unsafe { buf.align_to_mut::<u32>() };
let head = min(len, (4 - (ptr as usize & 3)) & 3); apply_mask_fallback(&mut prefix, mask);
let head = prefix.len() & 3;
let mask_u32 = if head > 0 { let mask_u32 = if head > 0 {
unsafe {
xor_mem(ptr, mask_u32, head);
ptr = ptr.add(head);
}
len -= head;
if cfg!(target_endian = "big") { if cfg!(target_endian = "big") {
mask_u32.rotate_left(8 * head as u32) mask_u32.rotate_left(8 * head as u32)
} else { } else {
@ -50,45 +35,15 @@ fn apply_mask_fast32(buf: &mut [u8], mask: [u8; 4]) {
} else { } else {
mask_u32 mask_u32
}; };
for word in words.iter_mut() {
if len > 0 { *word ^= mask_u32;
debug_assert_eq!(ptr as usize % 4, 0);
}
// Properly aligned middle of the data.
while len > 4 {
unsafe {
*(ptr as *mut u32) ^= mask_u32;
ptr = ptr.offset(4);
len -= 4;
}
}
// Possible last block.
if len > 0 {
unsafe {
xor_mem(ptr, mask_u32, len);
}
} }
} apply_mask_fallback(&mut suffix, mask_u32.to_ne_bytes());
#[inline]
// TODO: copy_nonoverlapping here compiles to call memcpy. While it is not so inefficient,
// it could be done better. The compiler does not see that len is limited to 3.
unsafe fn xor_mem(ptr: *mut u8, mask: u32, len: usize) {
#[allow(deprecated)]
let mut b: u32 = uninitialized();
#[allow(trivial_casts)]
copy_nonoverlapping(ptr, &mut b as *mut _ as *mut u8, len);
b ^= mask;
#[allow(trivial_casts)]
copy_nonoverlapping(&b as *const _ as *const u8, ptr, len);
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*;
use super::{apply_mask_fallback, apply_mask_fast32};
#[test] #[test]
fn test_apply_mask() { fn test_apply_mask() {
@ -98,26 +53,21 @@ mod tests {
0x12, 0x03, 0x12, 0x03,
]; ];
// Check masking with proper alignment. for data_len in 0..=unmasked.len() {
{ let unmasked = &unmasked[0..data_len];
let mut masked = unmasked.clone(); // Check masking with different alignment.
apply_mask_fallback(&mut masked, mask); for off in 0..=3 {
if unmasked.len() < off {
let mut masked_fast = unmasked.clone(); continue;
apply_mask_fast32(&mut masked_fast, mask); }
let mut masked = unmasked.to_vec();
assert_eq!(masked, masked_fast); apply_mask_fallback(&mut masked[off..], mask);
}
let mut masked_fast = unmasked.to_vec();
// Check masking without alignment. apply_mask_fast32(&mut masked_fast[off..], mask);
{
let mut masked = unmasked.clone(); assert_eq!(masked, masked_fast);
apply_mask_fallback(&mut masked[1..], mask); }
let mut masked_fast = unmasked.clone();
apply_mask_fast32(&mut masked_fast[1..], mask);
assert_eq!(masked, masked_fast);
} }
} }
} }

Loading…
Cancel
Save