proxmox-mirrors/proxmox
1
0
Fork 1
mirror of https://git.proxmox.com/git/proxmox synced 2025-08-04 17:14:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

proxmox/tools/websocket: improve mask_bytes and create_frame interface

Browse Source 
by using a Box<[u8]> instead of a vector (we do not need it)

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
This commit is contained in:
Dominik Csapak 2020-07-15 11:08:40 +02:00 committed by Wolfgang Bumiller
parent a3faec8322
commit 861a855cdf
1 changed files with 20 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
proxmox/src/tools/websocket.rs
View File

@ -44,18 +44,17 @@ impl OpCode {
}
}
fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Box<[u8]>) {
let mask = match mask {
Some([0,0,0,0]) | None => return data,
Some([0,0,0,0]) | None => return,
Some(mask) => mask,
};
if data.len() < 32 {
let mut_data = data.as_mut_slice();
for i in 0..mut_data.len() {
mut_data[i] ^= mask[i%4];
for i in 0..data.len() {
data[i] ^= mask[i%4];
}
return data;
return;
}
let mut newmask: u32 = u32::from_le_bytes(mask);
@ -75,8 +74,6 @@ fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
*s ^= newmask as u8;
newmask = newmask.rotate_right(8);
}
data
}
/// Can be used to create a complete WebSocket Frame.
@ -91,7 +88,7 @@ fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
/// # use std::io;
/// # fn main() -> io::Result<()> {
/// let data = vec![0,1,2,3,4];
/// let frame = create_frame(None, data, OpCode::Text)?;
/// let frame = create_frame(None, &data, OpCode::Text)?;
/// assert_eq!(frame, vec![0b10000001, 5, 0, 1, 2, 3, 4]);
/// # Ok(())
/// # }
@ -104,7 +101,7 @@ fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
/// # use std::io;
/// # fn main() -> io::Result<()> {
/// let data = vec![0,1,2,3,4];
/// let frame = create_frame(Some([0u8, 1u8, 2u8, 3u8]), data, OpCode::Text)?;
/// let frame = create_frame(Some([0u8, 1u8, 2u8, 3u8]), &data, OpCode::Text)?;
/// assert_eq!(frame, vec![0b10000001, 0b10000101, 0, 1, 2, 3, 0, 0, 0, 0, 4]);
/// # Ok(())
/// # }
@ -117,7 +114,7 @@ fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
/// # use std::io;
/// # fn main() -> io::Result<()> {
/// let data = vec![0,1,2,3,4];
/// let frame = create_frame(None, data, OpCode::Ping)?;
/// let frame = create_frame(None, &data, OpCode::Ping)?;
/// assert_eq!(frame, vec![0b10001001, 0b00000101, 0, 1, 2, 3, 4]);
/// # Ok(())
/// # }
@ -125,7 +122,7 @@ fn mask_bytes(mask: Option<[u8; 4]>, data: &mut Vec<u8>) -> &mut Vec<u8> {
/// ```
pub fn create_frame(
mask: Option<[u8; 4]>,
mut data: Vec<u8>,
data: &[u8],
frametype: OpCode,
) -> io::Result<Vec<u8>> {
let first_byte = 0b10000000 | (frametype as u8);
@ -155,8 +152,10 @@ pub fn create_frame(
if let Some(mask) = mask {
buf.extend_from_slice(&mask);
}
let mut data = data.to_vec().into_boxed_slice();
mask_bytes(mask, &mut data);
buf.append(&mut mask_bytes(mask, &mut data));
buf.append(&mut data.into_vec());
Ok(buf)
}
@ -212,7 +211,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for WebSocketWriter<W> {
if this.frame.is_none() {
// create frame buf
let frame = match create_frame(this.mask, buf.to_vec(), frametype) {
let frame = match create_frame(this.mask, buf, frametype) {
Ok(f) => f,
Err(e) => {
return Poll::Ready(Err(e));
@ -285,7 +284,7 @@ impl FrameHeader {
/// # use proxmox::tools::websocket::*;
/// # use std::io;
/// # fn main() -> io::Result<()> {
/// let frame = create_frame(None, vec![0,1,2,3], OpCode::Ping)?;
/// let frame = create_frame(None, &[0,1,2,3], OpCode::Ping)?;
/// let header = FrameHeader::try_from_bytes(&frame[..1])?;
/// match header {
/// Ok(_) => unreachable!(),
@ -514,12 +513,11 @@ impl<R: AsyncReadExt + Unpin + Send + 'static> AsyncRead for WebSocketReader<R>
if header.is_control_frame() {
if read_buffer.len() >= header.payload_len {
let mut data = read_buffer.remove_data(header.payload_len);
mask_bytes(header.mask, &mut data);
(this.callback)(
header.frametype,
mask_bytes(
header.mask,
&mut read_buffer.remove_data(header.payload_len).into_vec(),
),
&data,
);
this.state = if read_buffer.is_empty() {
ReaderState::NoData
@ -538,8 +536,9 @@ impl<R: AsyncReadExt + Unpin + Send + 'static> AsyncRead for WebSocketReader<R>
let len = min(buf.len() - offset, min(header.payload_len, read_buffer.len()));
let mut data = read_buffer.remove_data(len).into_vec();
buf[offset..offset+len].copy_from_slice(mask_bytes(header.mask, &mut data));
let mut data = read_buffer.remove_data(len);
mask_bytes(header.mask, &mut data);
buf[offset..offset+len].copy_from_slice(&data);
offset += len;
header.payload_len -= len;