use crate::{
credentials::Credentials,
packet::{control::Tag, stream, WireVersion},
};
use core::fmt;
use s2n_codec::{
decoder_invariant, CheckedRange, DecoderBufferMut, DecoderBufferMutResult as R, DecoderError,
};
use s2n_quic_core::{assume, frame::FrameMut, varint::VarInt};
type PacketNumber = VarInt;
pub trait Validator {
fn validate_tag(&mut self, tag: Tag) -> Result<(), DecoderError>;
}
impl Validator for () {
#[inline]
fn validate_tag(&mut self, _tag: Tag) -> Result<(), DecoderError> {
Ok(())
}
}
impl Validator for Tag {
#[inline]
fn validate_tag(&mut self, actual: Tag) -> Result<(), DecoderError> {
decoder_invariant!(*self == actual, "unexpected packet type");
Ok(())
}
}
impl<A, B> Validator for (A, B)
where
A: Validator,
B: Validator,
{
#[inline]
fn validate_tag(&mut self, tag: Tag) -> Result<(), DecoderError> {
self.0.validate_tag(tag)?;
self.1.validate_tag(tag)?;
Ok(())
}
}
pub struct Packet<'a> {
tag: Tag,
wire_version: WireVersion,
credentials: Credentials,
source_queue_id: Option<VarInt>,
stream_id: Option<stream::Id>,
packet_number: PacketNumber,
header: &'a mut [u8],
application_header: CheckedRange,
control_data: CheckedRange,
auth_tag: &'a mut [u8],
}
impl fmt::Debug for Packet<'_> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let header = &*self.header;
let mut s = f.debug_struct("control::Packet");
s.field("tag", &self.tag)
.field("wire_version", &self.wire_version)
.field("credentials", &self.credentials)
.field("source_queue_id", &self.source_queue_id)
.field("stream_id", &self.stream_id)
.field("packet_number", &self.packet_number);
if !self.application_header.is_empty() {
s.field("application_header", &self.application_header.get(header));
}
if !self.control_data.is_empty() {
s.field("control_data", &self.control_data.get(header));
}
s.field("auth_tag", &self.auth_tag).finish()
}
}
impl Packet<'_> {
#[inline]
pub fn tag(&self) -> Tag {
self.tag
}
#[inline]
pub fn wire_version(&self) -> WireVersion {
self.wire_version
}
#[inline]
pub fn credentials(&self) -> &Credentials {
&self.credentials
}
#[inline]
pub fn source_queue_id(&self) -> Option<VarInt> {
self.source_queue_id
}
#[inline]
pub fn stream_id(&self) -> Option<&stream::Id> {
self.stream_id.as_ref()
}
#[inline]
pub fn packet_number(&self) -> PacketNumber {
self.packet_number
}
#[inline]
pub fn application_header(&self) -> &[u8] {
self.application_header.get(self.header)
}
#[inline]
pub fn control_data(&self) -> &[u8] {
self.control_data.get(self.header)
}
#[inline]
pub fn control_data_mut(&mut self) -> &mut [u8] {
self.control_data.get_mut(self.header)
}
#[inline]
pub fn control_frames_mut(&mut self) -> ControlFramesMut<'_> {
ControlFramesMut {
buffer: self.control_data.get_mut(self.header),
}
}
#[inline]
pub fn header(&self) -> &[u8] {
self.header
}
#[inline]
pub fn auth_tag(&self) -> &[u8] {
self.auth_tag
}
#[inline]
pub fn total_len(&self) -> usize {
self.header.len() + self.auth_tag.len()
}
#[inline(always)]
pub fn decode<V: Validator>(
buffer: DecoderBufferMut,
mut validator: V,
crypto_tag_len: usize,
) -> R<Packet> {
let (
tag,
wire_version,
credentials,
source_queue_id,
stream_id,
packet_number,
header_len,
total_header_len,
application_header_len,
control_data_len,
) = {
let buffer = buffer.peek();
unsafe {
assume!(
crypto_tag_len >= 16,
"tag len needs to be at least 16 bytes"
);
}
let start_len = buffer.len();
let (tag, buffer) = buffer.decode()?;
validator.validate_tag(tag)?;
let (credentials, buffer) = buffer.decode()?;
let (wire_version, buffer) = buffer.decode()?;
let (stream_id, buffer) = if tag.is_stream() {
let (stream_id, buffer) = buffer.decode()?;
(Some(stream_id), buffer)
} else {
(None, buffer)
};
let (source_queue_id, buffer) = if tag.has_source_queue_id() {
let (v, buffer) = buffer.decode()?;
(Some(v), buffer)
} else {
(None, buffer)
};
let (packet_number, buffer) = buffer.decode::<VarInt>()?;
let (control_data_len, buffer) = buffer.decode::<VarInt>()?;
let (application_header_len, buffer) = if tag.has_application_header() {
let (application_header_len, buffer) = buffer.decode::<VarInt>()?;
((*application_header_len) as usize, buffer)
} else {
(0, buffer)
};
let header_len = start_len - buffer.len();
let buffer = buffer.skip(application_header_len)?;
let buffer = buffer.skip(*control_data_len as _)?;
let total_header_len = start_len - buffer.len();
let buffer = buffer.skip(crypto_tag_len)?;
let _ = buffer;
(
tag,
wire_version,
credentials,
source_queue_id,
stream_id,
packet_number,
header_len,
total_header_len,
application_header_len,
control_data_len,
)
};
unsafe {
assume!(buffer.len() >= total_header_len);
}
let (header, buffer) = buffer.decode_slice(total_header_len)?;
let (application_header, control_data) = {
let buffer = header.peek();
unsafe {
assume!(buffer.len() >= header_len);
}
let buffer = buffer.skip(header_len)?;
unsafe {
assume!(buffer.len() >= application_header_len);
}
let (application_header, buffer) =
buffer.skip_into_range(application_header_len, &header)?;
unsafe {
assume!(buffer.len() >= *control_data_len as usize);
}
let (control_data, _) = buffer.skip_into_range(*control_data_len as usize, &header)?;
(application_header, control_data)
};
let header = header.into_less_safe_slice();
let (auth_tag, buffer) = buffer.decode_slice(crypto_tag_len)?;
let auth_tag = auth_tag.into_less_safe_slice();
let packet = Packet {
tag,
wire_version,
credentials,
source_queue_id,
stream_id,
packet_number,
header,
application_header,
control_data,
auth_tag,
};
Ok((packet, buffer))
}
}
pub struct ControlFramesMut<'a> {
buffer: &'a mut [u8],
}
impl<'a> ControlFramesMut<'a> {
#[inline]
pub(crate) fn new(buffer: &'a mut [u8]) -> Self {
Self { buffer }
}
}
impl<'a> Iterator for ControlFramesMut<'a> {
type Item = Result<FrameMut<'a>, s2n_codec::DecoderError>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.buffer.is_empty() {
return None;
}
let buffer = unsafe {
core::mem::transmute::<&mut [u8], &mut [u8]>(self.buffer)
};
match DecoderBufferMut::new(buffer).decode::<FrameMut>() {
Ok((frame, remaining)) => {
self.buffer = remaining.into_less_safe_slice();
Some(Ok(frame))
}
Err(err) => {
self.buffer = &mut [];
Some(Err(err))
}
}
}
}