use std::collections::VecDeque;
use std::num::NonZeroUsize;
use o3::buffer::{ByteRing, ByteSink, CapacityError, Shared};
use crate::conn::{ConnError, Settings};
use crate::frame::{
self, Continuation, ErrorCode, Flags, FrameHeader, HEADER_LEN, Headers, RstStream,
WindowIncrement, WindowUpdate,
};
use crate::hpack;
use crate::stream::StreamId;
pub(crate) const COMMITTED_CONTROL_SLACK: usize = HEADER_LEN + 4;
#[derive(Clone, Copy)]
pub(crate) struct EgressCapacity {
logical: NonZeroUsize,
storage: NonZeroUsize,
}
impl EgressCapacity {
pub(crate) fn new(logical: usize) -> Option<Self> {
let logical = NonZeroUsize::new(logical)?;
let storage = logical
.get()
.checked_add(COMMITTED_CONTROL_SLACK)
.and_then(NonZeroUsize::new)?;
Some(Self { logical, storage })
}
pub(crate) const fn saturating(logical: NonZeroUsize) -> Self {
Self {
logical,
storage: logical.saturating_add(COMMITTED_CONTROL_SLACK),
}
}
}
struct SharedPayload {
prefix_len: usize,
body: Shared,
}
pub(crate) enum Write {
Buffered(usize),
Split { prefix_len: usize, body: Shared },
}
pub(crate) struct Egress {
bytes: ByteRing,
capacity: usize,
encoder: hpack::Encoder,
header_block: Vec<u8>,
payloads: VecDeque<SharedPayload>,
payload_bytes: usize,
unassigned_prefix: usize,
}
impl Egress {
pub(crate) fn from_capacity(
capacity: EgressCapacity,
table_size: usize,
header_capacity: usize,
) -> Self {
let EgressCapacity { logical, storage } = capacity;
Self {
bytes: ByteRing::with_capacity(storage),
capacity: logical.get(),
encoder: hpack::Encoder::new(table_size),
header_block: Vec::with_capacity(header_capacity),
payloads: VecDeque::new(),
payload_bytes: 0,
unassigned_prefix: 0,
}
}
pub(crate) fn raw_mut(&mut self) -> &mut Self {
self
}
pub(crate) fn first(&self) -> &[u8] {
let (first, _) = self.slices();
first
}
pub(crate) fn slices(&self) -> (&[u8], &[u8]) {
let limit = self
.payloads
.front()
.map_or(self.bytes.len(), |payload| payload.prefix_len);
let (first, second) = self.bytes.as_slices();
let first_len = first.len().min(limit);
let second_len = second.len().min(limit - first_len);
(&first[..first_len], &second[..second_len])
}
pub(crate) fn drain(&mut self, bytes: usize) {
let available = self
.payloads
.front()
.map_or(self.bytes.len(), |payload| payload.prefix_len);
let consumed = bytes.min(available);
self.bytes.consume_prefix_up_to(consumed);
if let Some(payload) = self.payloads.front_mut() {
payload.prefix_len -= consumed;
} else {
self.unassigned_prefix -= consumed;
}
}
pub(crate) fn drain_into(&mut self, write: &mut [u8]) -> usize {
let (first, second) = self.slices();
let first_len = first.len().min(write.len());
write[..first_len].copy_from_slice(&first[..first_len]);
let second_len = second.len().min(write.len() - first_len);
write[first_len..first_len + second_len].copy_from_slice(&second[..second_len]);
let written = first_len + second_len;
self.drain(written);
written
}
pub(crate) fn take_write(&mut self, write: &mut [u8]) -> Write {
let Some(prefix_len) = self.payloads.front().map(|payload| payload.prefix_len) else {
return Write::Buffered(self.drain_into(write));
};
let capacity = write.len().min(prefix_len);
let written = self.drain_into(&mut write[..capacity]);
if written != prefix_len {
return Write::Buffered(written);
}
let Some(payload) = self.payloads.pop_front() else {
return Write::Buffered(written);
};
self.payload_bytes -= payload.body.len();
Write::Split {
prefix_len: written,
body: payload.body,
}
}
pub(crate) fn len(&self) -> usize {
self.bytes.len() + self.payload_bytes
}
pub(crate) fn reserve(&self, additional: usize) -> Result<(), ConnError> {
if additional > self.bytes.remaining()
|| additional > self.capacity.saturating_sub(self.len())
{
Err(ConnError::Overload)
} else {
Ok(())
}
}
pub(crate) fn reserve_split(
&self,
prefix_len: usize,
payload_len: usize,
) -> Result<(), ConnError> {
let additional = prefix_len
.checked_add(payload_len)
.ok_or(ConnError::FrameSize)?;
if prefix_len > self.bytes.remaining()
|| additional > self.capacity.saturating_sub(self.len())
{
Err(ConnError::Overload)
} else {
Ok(())
}
}
pub(crate) fn over_capacity(&self) -> bool {
self.len() > self.capacity
}
pub(crate) fn queue_shared(&mut self, body: Shared) {
if body.is_empty() {
return;
}
self.payload_bytes += body.len();
self.payloads.push_back(SharedPayload {
prefix_len: std::mem::take(&mut self.unassigned_prefix),
body,
});
}
pub(crate) fn try_extend_from_slice(
&mut self,
bytes: &[u8],
) -> Result<(), o3::buffer::CapacityError> {
self.bytes.try_extend_from_slice(bytes)?;
self.unassigned_prefix += bytes.len();
Ok(())
}
pub(crate) fn initial_settings(&mut self, settings: &Settings) -> Result<(), ConnError> {
let count = settings.param_count();
let header = FrameHeader {
length: frame::FrameLength::from_usize(count * 6).ok_or(ConnError::FrameSize)?,
kind: frame::Type::Settings,
flags: Flags(0),
stream_id: StreamId::CONNECTION,
}
.wire_bytes();
let (params, len) = settings.wire_bytes();
self.write_slices([&header, ¶ms[..len]])?;
Ok(())
}
pub(crate) fn settings_ack(&mut self) -> Result<(), ConnError> {
self.reserve(HEADER_LEN)?;
FrameHeader {
length: frame::FrameLength::ZERO,
kind: frame::Type::Settings,
flags: Flags(Flags::ACK),
stream_id: StreamId::CONNECTION,
}
.encode(self)?;
Ok(())
}
pub(crate) fn window_update(
&mut self,
stream_id: StreamId,
increment: WindowIncrement,
) -> Result<(), ConnError> {
self.reserve(HEADER_LEN + 4)?;
WindowUpdate {
stream_id,
increment,
}
.encode(self)?;
Ok(())
}
pub(crate) fn data_control(
&mut self,
stream_id: StreamId,
connection_update: Option<u32>,
stream_update: Option<u32>,
reset: Option<ErrorCode>,
) -> Result<(), ConnError> {
let connection = Self::update_wire(StreamId::CONNECTION, connection_update)?;
let stream = Self::update_wire(stream_id, stream_update)?;
let reset = match reset {
Some(error) => Some(
RstStream::new(stream_id, error)
.ok_or(ConnError::Protocol)?
.wire_bytes(),
),
None => None,
};
let connection = connection.as_ref().map_or(&[][..], |wire| wire.as_slice());
let stream = stream.as_ref().map_or(&[][..], |wire| wire.as_slice());
let reset = reset.as_ref().map_or(&[][..], |wire| wire.as_slice());
let additional = connection.len() + stream.len() + reset.len();
if additional == 0 {
return Ok(());
}
self.reserve(additional)?;
self.write_slices([connection, stream, reset])?;
Ok(())
}
fn update_wire(
stream_id: StreamId,
increment: Option<u32>,
) -> Result<Option<[u8; HEADER_LEN + 4]>, ConnError> {
let Some(increment) = increment else {
return Ok(None);
};
Ok(Some(
WindowUpdate {
stream_id,
increment: WindowIncrement::new(increment).ok_or(ConnError::FlowControl)?,
}
.wire_bytes(),
))
}
pub(crate) fn committed_reset(
&mut self,
stream_id: StreamId,
error: ErrorCode,
) -> Result<(), ConnError> {
let wire = RstStream::new(stream_id, error)
.ok_or(ConnError::Protocol)?
.wire_bytes();
self.write_slices([&wire])?;
if self.over_capacity() {
Err(ConnError::Overload)
} else {
Ok(())
}
}
pub(crate) fn set_header_table_size(&mut self, size: usize) {
self.encoder.set_max_size(size);
}
pub(crate) fn headers<'a, I>(
&mut self,
stream_id: StreamId,
headers: I,
end_stream: bool,
max_frame: usize,
) -> Result<(), ConnError>
where
I: IntoIterator<Item = hpack::Header<'a>>,
{
let mut block = core::mem::take(&mut self.header_block);
block.clear();
self.encoder.encode(headers, &mut block);
let frames = block.len().max(1).div_ceil(max_frame);
let additional = frames
.checked_mul(HEADER_LEN)
.and_then(|heads| block.len().checked_add(heads))
.ok_or(ConnError::FrameSize);
let result = additional.and_then(|additional| self.reserve(additional));
if result.is_err() {
self.encoder.discard_block();
self.header_block = block;
return result;
}
if block.len() <= max_frame {
Headers::new(stream_id, end_stream, true, None, &block)
.ok_or(ConnError::FrameSize)?
.encode(self)?;
} else {
let (first, rest) = block.split_at(max_frame);
Headers::new(stream_id, end_stream, false, None, first)
.ok_or(ConnError::FrameSize)?
.encode(self)?;
let mut position = 0;
while position < rest.len() {
let take = (rest.len() - position).min(max_frame);
let end = position + take;
Continuation::new(stream_id, end == rest.len(), &rest[position..end])
.ok_or(ConnError::FrameSize)?
.encode(self)?;
position = end;
}
}
self.header_block = block;
Ok(())
}
}
impl ByteSink for Egress {
type Error = CapacityError;
fn write_slices<const N: usize>(&mut self, slices: [&[u8]; N]) -> Result<(), Self::Error> {
let before = self.bytes.len();
self.bytes.try_extend_from_slices(slices)?;
self.unassigned_prefix += self.bytes.len() - before;
Ok(())
}
}