use crate::{
collection::{ExpansionTy, Vector},
misc::{ConnectionState, Lease, LeaseMut},
rng::Rng,
stream::StreamWriter,
web_socket::{
Frame, FrameMut, OpCode, compression::NegotiatedCompression, misc::has_masked_frame,
unmask::unmask,
},
};
pub(crate) fn manage_compression<NC, P, const IS_CLIENT: bool>(
frame: &mut Frame<P, IS_CLIENT>,
nc_rsv1: u8,
) -> bool
where
NC: NegotiatedCompression,
P: Lease<[u8]>,
{
if NC::IS_NOOP {
return false;
}
let mut should_compress = false;
if !frame.op_code().is_control() {
let [first, _] = frame.header_first_two_mut();
should_compress = nc_rsv1 != 0;
*first |= nc_rsv1;
}
should_compress
}
pub(crate) fn manage_frame_compression<'cb, NC, P, R, const IS_CLIENT: bool>(
connection_state: &mut ConnectionState,
nc: &mut NC,
nc_rsv1: u8,
frame: &mut Frame<P, IS_CLIENT>,
no_masking: bool,
rng: &mut R,
writer_buffer: &'cb mut Vector<u8>,
) -> crate::Result<FrameMut<'cb, IS_CLIENT>>
where
NC: NegotiatedCompression,
P: LeaseMut<[u8]>,
R: Rng,
{
if frame.op_code() == OpCode::Close {
*connection_state = ConnectionState::Closed;
}
let mut compressed_frame = compress_frame(frame, nc, nc_rsv1, writer_buffer)?;
mask_frame(&mut compressed_frame, no_masking, rng);
Ok(compressed_frame)
}
pub(crate) fn manage_normal_frame<P, R, const IS_CLIENT: bool>(
connection_state: &mut ConnectionState,
frame: &mut Frame<P, IS_CLIENT>,
no_masking: bool,
rng: &mut R,
) where
P: LeaseMut<[u8]>,
R: Rng,
{
if frame.op_code() == OpCode::Close {
*connection_state = ConnectionState::Closed;
}
mask_frame(frame, no_masking, rng);
}
pub(crate) async fn write_frame<NC, P, R, SW, const IS_CLIENT: bool>(
connection_state: &mut ConnectionState,
frame: &mut Frame<P, IS_CLIENT>,
no_masking: bool,
nc: &mut NC,
nc_rsv1: u8,
rng: &mut R,
stream_writer: &mut SW,
writer_buffer: &mut Vector<u8>,
) -> crate::Result<()>
where
NC: NegotiatedCompression,
P: LeaseMut<[u8]>,
R: Rng,
SW: StreamWriter,
{
if manage_compression::<NC, _, IS_CLIENT>(frame, nc_rsv1) {
let fr = manage_frame_compression(
connection_state,
nc,
nc_rsv1,
frame,
no_masking,
rng,
writer_buffer,
)?;
stream_writer.write_all_vectored(&[fr.header(), fr.payload()]).await?;
} else {
manage_normal_frame::<_, _, IS_CLIENT>(connection_state, frame, no_masking, rng);
let (header, payload) = frame.header_and_payload_mut();
stream_writer.write_all_vectored(&[header, payload.lease()]).await?;
}
Ok(())
}
fn compress_frame<'cb, P, NC, const IS_CLIENT: bool>(
frame: &mut Frame<P, IS_CLIENT>,
nc: &mut NC,
nc_rsv1: u8,
writer_buffer: &'cb mut Vector<u8>,
) -> crate::Result<FrameMut<'cb, IS_CLIENT>>
where
P: LeaseMut<[u8]>,
NC: NegotiatedCompression,
{
let additional = frame.payload().lease().len().wrapping_add(128);
writer_buffer.clear();
let mut payload_len = nc.compress(
frame.payload().lease(),
writer_buffer,
|local_writer_buffer| {
local_writer_buffer.expand(ExpansionTy::Additional(additional), 0)?;
Ok(local_writer_buffer)
},
|local_writer_buffer, written| {
local_writer_buffer.expand(ExpansionTy::Additional(additional), 0)?;
Ok(local_writer_buffer.get_mut(written..).unwrap_or_default())
},
)?;
if frame.fin() {
payload_len = payload_len.saturating_sub(4);
}
Ok(FrameMut::new(
frame.fin(),
frame.op_code(),
writer_buffer.get_mut(..payload_len).unwrap_or_default(),
nc_rsv1,
))
}
fn mask_frame<P, R, const IS_CLIENT: bool>(
frame: &mut Frame<P, IS_CLIENT>,
no_masking: bool,
rng: &mut R,
) where
P: LeaseMut<[u8]>,
R: Rng,
{
if IS_CLIENT && !no_masking && !has_masked_frame(*frame.header_first_two_mut()[1]) {
let mask: [u8; 4] = rng.u8_4();
frame.set_mask(mask);
unmask(frame.payload_mut().lease_mut(), mask);
}
}