use bytes::Bytes;
use super::{
ack::ack_frame_with_ecn, connection_close::connection_close_frame_at_layer,
crypto::be_crypto_frame, data_blocked::be_data_blocked_frame,
datagram::datagram_frame_with_flag, max_data::be_max_data_frame,
max_stream_data::be_max_stream_data_frame, max_streams::max_streams_frame_with_dir,
new_connection_id::be_new_connection_id_frame, new_token::be_new_token_frame,
path_challenge::be_path_challenge_frame, path_response::be_path_response_frame,
reset_stream::be_reset_stream_frame, retire_connection_id::be_retire_connection_id_frame,
stop_sending::be_stop_sending_frame, stream::stream_frame_with_flag,
stream_data_blocked::be_stream_data_blocked_frame,
streams_blocked::streams_blocked_frame_with_dir, *,
};
use crate::util::ContinuousData;
fn complete_frame(
frame_type: FrameType,
raw: Bytes,
) -> impl Fn(&[u8]) -> nom::IResult<&[u8], Frame> {
use nom::{Parser, combinator::map};
move |input: &[u8]| match frame_type {
FrameType::Padding => Ok((input, Frame::Padding(PaddingFrame))),
FrameType::Ping => Ok((input, Frame::Ping(PingFrame))),
FrameType::ConnectionClose(layer) => {
map(connection_close_frame_at_layer(layer), Frame::Close).parse(input)
}
FrameType::NewConnectionId => {
map(be_new_connection_id_frame, Frame::NewConnectionId).parse(input)
}
FrameType::RetireConnectionId => {
map(be_retire_connection_id_frame, Frame::RetireConnectionId).parse(input)
}
FrameType::DataBlocked => map(be_data_blocked_frame, Frame::DataBlocked).parse(input),
FrameType::MaxData => map(be_max_data_frame, Frame::MaxData).parse(input),
FrameType::PathChallenge => map(be_path_challenge_frame, Frame::Challenge).parse(input),
FrameType::PathResponse => map(be_path_response_frame, Frame::Response).parse(input),
FrameType::HandshakeDone => Ok((input, Frame::HandshakeDone(HandshakeDoneFrame))),
FrameType::NewToken => map(be_new_token_frame, Frame::NewToken).parse(input),
FrameType::Ack(ecn) => map(ack_frame_with_ecn(ecn), Frame::Ack).parse(input),
FrameType::ResetStream => {
map(be_reset_stream_frame, |f| Frame::StreamCtl(f.into())).parse(input)
}
FrameType::StopSending => {
map(be_stop_sending_frame, |f| Frame::StreamCtl(f.into())).parse(input)
}
FrameType::MaxStreamData => {
map(be_max_stream_data_frame, |f| Frame::StreamCtl(f.into())).parse(input)
}
FrameType::MaxStreams(dir) => map(max_streams_frame_with_dir(dir), |f| {
Frame::StreamCtl(f.into())
})
.parse(input),
FrameType::StreamsBlocked(dir) => map(streams_blocked_frame_with_dir(dir), |f| {
Frame::StreamCtl(f.into())
})
.parse(input),
FrameType::StreamDataBlocked => {
map(be_stream_data_blocked_frame, |f| Frame::StreamCtl(f.into())).parse(input)
}
FrameType::Crypto => {
let (input, frame) = be_crypto_frame(input)?;
let start = raw.len() - input.len();
let len = frame.length() as usize;
if input.len() < len {
Err(nom::Err::Incomplete(nom::Needed::new(len - input.len())))
} else {
let data = raw.slice(start..start + len);
Ok((&input[len..], Frame::Crypto(frame, data)))
}
}
FrameType::Stream(flags) => {
let (input, frame) = stream_frame_with_flag(flags)(input)?;
let start = raw.len() - input.len();
let len = frame.len();
if input.len() < len {
Err(nom::Err::Incomplete(nom::Needed::new(len - input.len())))
} else {
let data = raw.slice(start..start + len);
Ok((&input[len..], Frame::Stream(frame, data)))
}
}
FrameType::Datagram(with_len) => {
let (input, frame) = datagram_frame_with_flag(with_len)(input)?;
let start = raw.len() - input.len();
match frame.encode_len() {
true if frame.len().into_inner() > input.len() as u64 => Err(nom::Err::Incomplete(
nom::Needed::new((frame.len().into_inner() - input.len() as u64) as usize),
)),
true => {
let data = raw.slice(start..start + frame.len().into_inner() as usize);
Ok((
&input[frame.len().into_inner() as usize..],
Frame::Datagram(frame, data),
))
}
false => {
let data = raw.slice(start..);
Ok((&[], Frame::Datagram(frame, data)))
}
}
}
}
}
pub fn be_frame(raw: &Bytes, packet_type: Type) -> Result<(usize, Frame, FrameType), Error> {
let input = raw.as_ref();
let (remain, frame_type) = be_frame_type(input)?;
if !frame_type.belongs_to(packet_type) {
return Err(Error::WrongType(frame_type, packet_type));
}
let (remain, frame) = complete_frame(frame_type, raw.clone())(remain).map_err(|e| match e {
ne @ nom::Err::Incomplete(_) => {
nom::Err::Error(Error::IncompleteFrame(frame_type, ne.to_string()))
}
nom::Err::Error(ne) => {
nom::Err::Error(Error::ParseError(
frame_type,
ne.code.description().to_owned(),
))
}
_ => unreachable!("parsing frame never fails"),
})?;
Ok((input.len() - remain.len(), frame, frame_type))
}
pub trait WriteFrame<F>: bytes::BufMut {
fn put_frame(&mut self, frame: &F);
}
impl<B: BufMut, D: ContinuousData> WriteFrame<Frame<D>> for B
where
D: ContinuousData,
B: BufMut + ?Sized,
for<'b> &'b mut B: crate::util::WriteData<D>,
{
fn put_frame(&mut self, frame: &Frame<D>) {
let mut buf = self;
match frame {
Frame::Padding(f) => <&mut B as WriteFrame<PaddingFrame>>::put_frame(&mut buf, f),
Frame::Ping(f) => <&mut B as WriteFrame<PingFrame>>::put_frame(&mut buf, f),
Frame::Ack(f) => <&mut B as WriteFrame<AckFrame>>::put_frame(&mut buf, f),
Frame::Close(f) => <&mut B as WriteFrame<ConnectionCloseFrame>>::put_frame(&mut buf, f),
Frame::NewToken(f) => <&mut B as WriteFrame<NewTokenFrame>>::put_frame(&mut buf, f),
Frame::MaxData(f) => <&mut B as WriteFrame<MaxDataFrame>>::put_frame(&mut buf, f),
Frame::DataBlocked(f) => {
<&mut B as WriteFrame<DataBlockedFrame>>::put_frame(&mut buf, f)
}
Frame::NewConnectionId(f) => {
<&mut B as WriteFrame<NewConnectionIdFrame>>::put_frame(&mut buf, f)
}
Frame::RetireConnectionId(f) => {
<&mut B as WriteFrame<RetireConnectionIdFrame>>::put_frame(&mut buf, f)
}
Frame::HandshakeDone(f) => {
<&mut B as WriteFrame<HandshakeDoneFrame>>::put_frame(&mut buf, f)
}
Frame::Challenge(f) => {
<&mut B as WriteFrame<PathChallengeFrame>>::put_frame(&mut buf, f)
}
Frame::Response(f) => <&mut B as WriteFrame<PathResponseFrame>>::put_frame(&mut buf, f),
Frame::StreamCtl(f) => <&mut B as WriteFrame<StreamCtlFrame>>::put_frame(&mut buf, f),
Frame::Stream(f, d) => buf.put_data_frame(f, d),
Frame::Crypto(f, d) => buf.put_data_frame(f, d),
Frame::Datagram(f, d) => buf.put_data_frame(f, d),
}
}
}
pub trait WriteDataFrame<F, D: ContinuousData>: bytes::BufMut {
fn put_data_frame(&mut self, frame: &F, data: &D);
}