mercuryworkshop-mirror/epoxy-tls
1
0
Fork 0
mirror of https://github.com/MercuryWorkshop/epoxy-tls.git synced 2025-05-13 14:30:02 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

read split frames

Browse source
This commit is contained in:
Toshit Chawda 2024-07-22 11:04:12 -07:00
parent 7f37c8338e
commit 76eeec87dc
No known key found for this signature in database
GPG key ID: 91480ED99E2B3D9D
5 changed files with 653 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

106
wisp/src/fastwebsockets.rs
View file

@ -3,7 +3,8 @@ use std::ops::Deref;
use async_trait::async_trait;
use bytes::BytesMut;
use fastwebsockets::{
CloseCode, FragmentCollectorRead, Frame, OpCode, Payload, WebSocketError, WebSocketWrite,
CloseCode, FragmentCollectorRead, Frame, OpCode, Payload, WebSocketError, WebSocketRead,
WebSocketWrite,
};
use tokio::io::{AsyncRead, AsyncWrite};
@ -88,13 +89,114 @@ impl<S: AsyncRead + Unpin + Send> crate::ws::WebSocketRead for FragmentCollector
}
}
#[async_trait]
impl<S: AsyncRead + Unpin + Send> crate::ws::WebSocketRead for WebSocketRead<S> {
async fn wisp_read_frame(
&mut self,
tx: &LockedWebSocketWrite,
) -> Result<crate::ws::Frame<'static>, WispError> {
let mut frame = self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
.await?;
if frame.opcode == OpCode::Continuation {
return Err(WispError::WsImplError(Box::new(
WebSocketError::InvalidContinuationFrame,
)));
}
let mut buf = BytesMut::from(frame.payload);
let opcode = frame.opcode;
while !frame.fin {
frame = self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
.await?;
if frame.opcode != OpCode::Continuation {
return Err(WispError::WsImplError(Box::new(
WebSocketError::InvalidContinuationFrame,
)));
}
buf.extend_from_slice(&frame.payload);
}
Ok(crate::ws::Frame {
opcode: opcode.into(),
payload: crate::ws::Payload::Bytes(buf),
finished: frame.fin,
})
}
async fn wisp_read_split(
&mut self,
tx: &LockedWebSocketWrite,
) -> Result<(crate::ws::Frame<'static>, Option<crate::ws::Frame<'static>>), WispError> {
let mut frame_cnt = 1;
let mut frame = self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
.await?;
let mut extra_frame = None;
if frame.opcode == OpCode::Continuation {
return Err(WispError::WsImplError(Box::new(
WebSocketError::InvalidContinuationFrame,
)));
}
let mut buf = BytesMut::from(frame.payload);
let opcode = frame.opcode;
while !frame.fin {
frame = self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
.await?;
if frame.opcode != OpCode::Continuation {
return Err(WispError::WsImplError(Box::new(
WebSocketError::InvalidContinuationFrame,
)));
}
if frame_cnt == 1 {
let payload = BytesMut::from(frame.payload);
extra_frame = Some(crate::ws::Frame {
opcode: opcode.into(),
payload: crate::ws::Payload::Bytes(payload),
finished: true,
});
} else if frame_cnt == 2 {
let extra_payload = extra_frame.take().unwrap().payload;
buf.extend_from_slice(&extra_payload);
buf.extend_from_slice(&frame.payload);
} else {
buf.extend_from_slice(&frame.payload);
}
frame_cnt += 1;
}
Ok((
crate::ws::Frame {
opcode: opcode.into(),
payload: crate::ws::Payload::Bytes(buf),
finished: frame.fin,
},
extra_frame,
))
}
}
#[async_trait]
impl<S: AsyncWrite + Unpin + Send> crate::ws::WebSocketWrite for WebSocketWrite<S> {
async fn wisp_write_frame(&mut self, frame: crate::ws::Frame<'_>) -> Result<(), WispError> {
self.write_frame(frame.into()).await.map_err(|e| e.into())
}
async fn wisp_write_split(&mut self, header: crate::ws::Frame<'_>, body: crate::ws::Frame<'_>) -> Result<(), WispError> {
async fn wisp_write_split(
&mut self,
header: crate::ws::Frame<'_>,
body: crate::ws::Frame<'_>,
) -> Result<(), WispError> {
let mut header = Frame::from(header);
header.fin = false;
self.write_frame(header).await?;