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