Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

#729: proto: write outgoing packets to caller-supplied memory #1697

Merged
merged 9 commits into from
Nov 1, 2023
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 32 additions & 20 deletions quinn-proto/src/connection/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -455,7 +455,12 @@ impl Connection {
/// `max_datagrams` specifies how many datagrams can be returned inside a
/// single Transmit using GSO. This must be at least 1.
#[must_use]
pub fn poll_transmit(&mut self, now: Instant, max_datagrams: usize) -> Option<Transmit> {
pub fn poll_transmit(
&mut self,
now: Instant,
max_datagrams: usize,
buf: &mut BytesMut,
) -> Option<Transmit> {
assert!(max_datagrams != 0);
let max_datagrams = match self.config.enable_segmentation_offload {
false => 1,
Expand All @@ -477,13 +482,14 @@ impl Connection {
SpaceId::Data,
"PATH_CHALLENGE queued without 1-RTT keys"
);
let mut buf = BytesMut::with_capacity(self.path.current_mtu() as usize);
let buf_capacity = self.path.current_mtu() as usize;
buf.reserve(self.path.current_mtu() as usize);

let buf_capacity = buf.capacity();

let mut builder = PacketBuilder::new(
now,
SpaceId::Data,
&mut buf,
buf,
buf_capacity,
0,
false,
Expand All @@ -501,13 +507,13 @@ impl Connection {
// sending a datagram of this size
builder.pad_to(MIN_INITIAL_SIZE);

builder.finish(self, &mut buf);
builder.finish(self, buf);
self.stats.udp_tx.datagrams += 1;
self.stats.udp_tx.ios += 1;
self.stats.udp_tx.bytes += buf.len() as u64;
return Some(Transmit {
destination,
contents: buf.freeze(),
size: buf.len(),
ecn: None,
segment_size: None,
src_ip: self.local_ip,
Expand Down Expand Up @@ -549,7 +555,6 @@ impl Connection {
&& self.peer_supports_ack_frequency();
}

let mut buf = BytesMut::new();
// Reserving capacity can provide more capacity than we asked for.
// However we are not allowed to write more than MTU size. Therefore
// the maximum capacity is tracked separately.
Expand Down Expand Up @@ -664,7 +669,7 @@ impl Connection {
// which will always send the maximum PDU.
builder.pad_to(self.path.current_mtu());

builder.finish_and_track(now, self, sent_frames.take(), &mut buf);
builder.finish_and_track(now, self, sent_frames.take(), buf);

debug_assert_eq!(buf.len(), buf_capacity, "Packet must be padded");
}
Expand All @@ -680,7 +685,7 @@ impl Connection {
// (e.g. purely containing ACKs), modern memory allocators
// (e.g. mimalloc and jemalloc) will pool certain allocation sizes
// and therefore this is still rather efficient.
buf.reserve(max_datagrams * self.path.current_mtu() as usize - buf.capacity());
buf.reserve(max_datagrams * self.path.current_mtu() as usize);
}
num_datagrams += 1;
coalesce = true;
Expand All @@ -690,7 +695,7 @@ impl Connection {
// datagram.
// Finish current packet without adding extra padding
if let Some(builder) = builder.take() {
builder.finish_and_track(now, self, sent_frames.take(), &mut buf);
builder.finish_and_track(now, self, sent_frames.take(), buf);
}
}

Expand Down Expand Up @@ -723,7 +728,7 @@ impl Connection {
let builder = builder.get_or_insert(PacketBuilder::new(
now,
space_id,
&mut buf,
buf,
buf_capacity,
(num_datagrams - 1) * (self.path.current_mtu() as usize),
ack_eliciting,
Expand All @@ -748,7 +753,7 @@ impl Connection {
self.receiving_ecn,
&mut SentFrames::default(),
&mut self.spaces[space_id],
&mut buf,
buf,
&mut self.stats,
);
}
Expand All @@ -764,22 +769,22 @@ impl Connection {
match self.state {
State::Closed(state::Closed { ref reason }) => {
if space_id == SpaceId::Data {
reason.encode(&mut buf, builder.max_size)
reason.encode(buf, builder.max_size)
} else {
frame::ConnectionClose {
error_code: TransportErrorCode::APPLICATION_ERROR,
frame_type: None,
reason: Bytes::new(),
}
.encode(&mut buf, builder.max_size)
.encode(buf, builder.max_size)
}
}
State::Draining => frame::ConnectionClose {
error_code: TransportErrorCode::NO_ERROR,
frame_type: None,
reason: Bytes::new(),
}
.encode(&mut buf, builder.max_size),
.encode(buf, builder.max_size),
_ => unreachable!(
"tried to make a close packet when the connection wasn't closed"
),
Expand All @@ -794,7 +799,7 @@ impl Connection {
let sent = self.populate_packet(
now,
space_id,
&mut buf,
buf,
buf_capacity - builder.tag_len,
builder.exact_number,
);
Expand Down Expand Up @@ -831,7 +836,7 @@ impl Connection {
builder.pad_to(MIN_INITIAL_SIZE);
}
let last_packet_number = builder.exact_number;
builder.finish_and_track(now, self, sent_frames, &mut buf);
builder.finish_and_track(now, self, sent_frames, buf);
self.path
.congestion
.on_sent(now, buf.len() as u64, last_packet_number);
Expand All @@ -857,7 +862,7 @@ impl Connection {
let mut builder = PacketBuilder::new(
now,
space_id,
&mut buf,
buf,
buf_capacity,
0,
true,
Expand All @@ -880,7 +885,7 @@ impl Connection {
non_retransmits: true,
..Default::default()
};
builder.finish_and_track(now, self, Some(sent_frames), &mut buf);
builder.finish_and_track(now, self, Some(sent_frames), buf);

self.stats.path.sent_plpmtud_probes += 1;
num_datagrams = 1;
Expand All @@ -901,7 +906,7 @@ impl Connection {

Some(Transmit {
destination: self.path.remote,
contents: buf.freeze(),
size: buf.len(),
ecn: if self.path.sending_ecn {
Some(EcnCodepoint::Ect0)
} else {
Expand Down Expand Up @@ -3415,6 +3420,13 @@ impl Connection {
self.state = State::Drained;
self.endpoint_events.push_back(EndpointEventInner::Drained);
}

/// Storage size required for the largest packet known to be supported by the current path
///
/// Buffers passed to [`Connection::poll_transmit`] should be at least this large.
pub fn current_mtu(&self) -> u16 {
self.path.current_mtu()
}
}

impl fmt::Debug for Connection {
Expand Down
45 changes: 22 additions & 23 deletions quinn-proto/src/endpoint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ impl Endpoint {
local_ip: Option<IpAddr>,
ecn: Option<EcnCodepoint>,
data: BytesMut,
buf: &mut BytesMut,
) -> Option<DatagramEvent> {
let datagram_len = data.len();
let (first_decode, remaining) = match PartialDecode::new(
Expand All @@ -147,13 +148,12 @@ impl Endpoint {
}
trace!("sending version negotiation");
// Negotiate versions
let mut buf = BytesMut::new();
Header::VersionNegotiate {
random: self.rng.gen::<u8>() | 0x40,
src_cid: dst_cid,
dst_cid: src_cid,
}
.encode(&mut buf);
.encode(buf);
// Grease with a reserved version
if version != 0x0a1a_2a3a {
buf.write::<u32>(0x0a1a_2a3a);
Expand All @@ -166,7 +166,7 @@ impl Endpoint {
return Some(DatagramEvent::Response(Transmit {
destination: remote,
ecn: None,
contents: buf.freeze(),
size: buf.len(),
segment_size: None,
src_ip: local_ip,
}));
Expand Down Expand Up @@ -205,7 +205,7 @@ impl Endpoint {
None => {
debug!("packet for unrecognized connection {}", dst_cid);
return self
.stateless_reset(datagram_len, addresses, dst_cid)
.stateless_reset(datagram_len, addresses, dst_cid, buf)
.map(DatagramEvent::Response);
}
};
Expand Down Expand Up @@ -233,7 +233,7 @@ impl Endpoint {
};
return match first_decode.finish(Some(&*crypto.header.remote)) {
Ok(packet) => {
self.handle_first_packet(now, addresses, ecn, packet, remaining, &crypto)
self.handle_first_packet(now, addresses, ecn, packet, remaining, &crypto, buf)
}
Err(e) => {
trace!("unable to decode initial packet: {}", e);
Expand All @@ -254,7 +254,7 @@ impl Endpoint {
//
if !dst_cid.is_empty() {
return self
.stateless_reset(datagram_len, addresses, dst_cid)
.stateless_reset(datagram_len, addresses, dst_cid, buf)
.map(DatagramEvent::Response);
}

Expand All @@ -267,6 +267,7 @@ impl Endpoint {
inciting_dgram_len: usize,
addresses: FourTuple,
dst_cid: &ConnectionId,
buf: &mut BytesMut,
) -> Option<Transmit> {
/// Minimum amount of padding for the stateless reset to look like a short-header packet
const MIN_PADDING_LEN: usize = 5;
Expand All @@ -285,7 +286,6 @@ impl Endpoint {
"sending stateless reset for {} to {}",
dst_cid, addresses.remote
);
let mut buf = BytesMut::new();
// Resets with at least this much padding can't possibly be distinguished from real packets
const IDEAL_MIN_PADDING_LEN: usize = MIN_PADDING_LEN + MAX_CID_SIZE;
let padding_len = if max_padding_len <= IDEAL_MIN_PADDING_LEN {
Expand All @@ -304,7 +304,7 @@ impl Endpoint {
Some(Transmit {
destination: addresses.remote,
ecn: None,
contents: buf.freeze(),
size: buf.len(),
segment_size: None,
src_ip: addresses.local_ip,
})
Expand Down Expand Up @@ -404,6 +404,7 @@ impl Endpoint {
mut packet: Packet,
rest: Option<BytesMut>,
crypto: &Keys,
buf: &mut BytesMut,
) -> Option<DatagramEvent> {
let (src_cid, dst_cid, token, packet_number, version) = match packet.header {
Header::Initial {
Expand Down Expand Up @@ -444,6 +445,7 @@ impl Endpoint {
crypto,
&src_cid,
TransportError::CONNECTION_REFUSED(""),
buf,
)));
}

Expand All @@ -460,6 +462,7 @@ impl Endpoint {
crypto,
&src_cid,
TransportError::PROTOCOL_VIOLATION("invalid destination CID length"),
buf,
)));
}

Expand Down Expand Up @@ -488,16 +491,15 @@ impl Endpoint {
version,
};

let mut buf = BytesMut::new();
let encode = header.encode(&mut buf);
let encode = header.encode(buf);
buf.put_slice(&token);
buf.extend_from_slice(&server_config.crypto.retry_tag(version, &dst_cid, &buf));
encode.finish(&mut buf, &*crypto.header.local, None);
buf.extend_from_slice(&server_config.crypto.retry_tag(version, &dst_cid, buf));
encode.finish(buf, &*crypto.header.local, None);

return Some(DatagramEvent::Response(Transmit {
destination: addresses.remote,
ecn: None,
contents: buf.freeze(),
size: buf.len(),
segment_size: None,
src_ip: addresses.local_ip,
}));
Expand All @@ -522,6 +524,7 @@ impl Endpoint {
crypto,
&src_cid,
TransportError::INVALID_TOKEN(""),
buf,
)));
}
}
Expand Down Expand Up @@ -569,7 +572,7 @@ impl Endpoint {
self.handle_event(ch, EndpointEvent(EndpointEventInner::Drained));
match e {
ConnectionError::TransportError(e) => Some(DatagramEvent::Response(
self.initial_close(version, addresses, crypto, &src_cid, e),
self.initial_close(version, addresses, crypto, &src_cid, e, buf),
)),
_ => None,
}
Expand Down Expand Up @@ -630,6 +633,7 @@ impl Endpoint {
crypto: &Keys,
remote_id: &ConnectionId,
reason: TransportError,
buf: &mut BytesMut,
) -> Transmit {
// We don't need to worry about CID collisions in initial closes because the peer
// shouldn't respond, and if it does, and the CID collides, we'll just drop the
Expand All @@ -644,21 +648,16 @@ impl Endpoint {
version,
};

let mut buf = BytesMut::new();
let partial_encode = header.encode(&mut buf);
let partial_encode = header.encode(buf);
let max_len =
INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
frame::Close::from(reason).encode(&mut buf, max_len);
frame::Close::from(reason).encode(buf, max_len);
buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
partial_encode.finish(
&mut buf,
&*crypto.header.local,
Some((0, &*crypto.packet.local)),
);
partial_encode.finish(buf, &*crypto.header.local, Some((0, &*crypto.packet.local)));
Transmit {
destination: addresses.remote,
ecn: None,
contents: buf.freeze(),
size: buf.len(),
segment_size: None,
src_ip: addresses.local_ip,
}
Expand Down
5 changes: 2 additions & 3 deletions quinn-proto/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@ mod tests;
pub mod transport_parameters;
mod varint;

use bytes::Bytes;
pub use varint::{VarInt, VarIntBoundsExceeded};

mod connection;
Expand Down Expand Up @@ -281,8 +280,8 @@ pub struct Transmit {
pub destination: SocketAddr,
/// Explicit congestion notification bits to set on the packet
pub ecn: Option<EcnCodepoint>,
/// Contents of the datagram
pub contents: Bytes,
/// Amount of data written to the caller-supplied buffer
pub size: usize,
/// The segment size if this transmission contains multiple datagrams.
/// This is `None` if the transmit only contains a single datagram
pub segment_size: Option<usize>,
Expand Down
Loading
Loading