1use std::{
30 pin::Pin,
31 task::{Context, Poll},
32};
33
34use bytes::{Buf, Bytes};
35use futures_util::{ready, stream, Stream, StreamExt};
36use tokio_util::sync::ReusableBoxFuture;
37
38use crate::h3::error::TransportError;
39use crate::h3::transport::{
40 self, Accept, BidiStream as BidiStreamTrait, OpenStreams, RecvStream as RecvStreamTrait,
41 SendStream as SendStreamTrait, UniStream as UniStreamTrait,
42};
43
44type BoxStream<T> = Pin<Box<dyn Stream<Item = T> + std::marker::Send + 'static>>;
46
47type AcceptBiItem = Result<(quinn::SendStream, quinn::RecvStream), quinn::ConnectionError>;
48type AcceptUniItem = Result<quinn::RecvStream, quinn::ConnectionError>;
49type OpenUniItem = Result<quinn::SendStream, quinn::ConnectionError>;
50
51pub struct Connection {
56 conn: quinn::Connection,
57 accept_bi: BoxStream<AcceptBiItem>,
64 accept_uni: BoxStream<AcceptUniItem>,
65 open_uni: Option<BoxStream<OpenUniItem>>,
67}
68
69impl Connection {
70 #[inline]
76 pub fn new(conn: quinn::Connection) -> Self {
77 let accept_bi = Box::pin(stream::unfold(conn.clone(), |conn| async move {
78 Some((conn.accept_bi().await, conn))
79 }));
80 let accept_uni = Box::pin(stream::unfold(conn.clone(), |conn| async move {
81 Some((conn.accept_uni().await, conn))
82 }));
83 Self {
84 conn,
85 accept_bi,
86 accept_uni,
87 open_uni: None,
88 }
89 }
90
91 #[inline]
97 pub fn close_reason(&self) -> Option<quinn::ConnectionError> {
98 self.conn.close_reason()
99 }
100}
101
102impl OpenStreams for Connection {
103 #[inline]
104 fn poll_open_uni(
105 &mut self,
106 cx: &mut Context<'_>,
107 ) -> Poll<Result<Box<dyn UniStreamTrait>, TransportError>> {
108 let stream = self.open_uni.get_or_insert_with(|| {
109 let conn = self.conn.clone();
110 Box::pin(stream::unfold(conn, |conn| async move {
111 Some((conn.open_uni().await, conn))
112 }))
113 });
114 match ready!(stream.poll_next_unpin(cx)) {
115 Some(Ok(stream)) => Poll::Ready(Ok(Box::new(UniStream::Send(Send::new(stream))))),
116 Some(Err(err)) => Poll::Ready(Err(map_connection_error(err))),
117 None => unreachable!("unfold stream never ends"),
118 }
119 }
120}
121
122impl Accept for Connection {
123 #[inline]
124 fn poll_accept(
125 &mut self,
126 cx: &mut Context<'_>,
127 ) -> Poll<Result<Option<Box<dyn BidiStreamTrait>>, TransportError>> {
128 match ready!(self.accept_bi.as_mut().poll_next_unpin(cx)) {
129 Some(Ok((send, recv))) => Poll::Ready(Ok(Some(Box::new(BidiStream {
130 send: Send::new(send),
131 recv: Recv::new(recv),
132 })))),
133 Some(Err(_)) => Poll::Ready(Ok(None)),
137 None => unreachable!("unfold stream never ends"),
138 }
139 }
140
141 #[inline]
142 fn poll_accept_uni(
143 &mut self,
144 cx: &mut Context<'_>,
145 ) -> Poll<Result<Option<Box<dyn UniStreamTrait>>, TransportError>> {
146 match ready!(self.accept_uni.as_mut().poll_next_unpin(cx)) {
147 Some(Ok(stream)) => Poll::Ready(Ok(Some(Box::new(UniStream::Recv(Recv::new(stream)))))),
148 Some(Err(_)) => Poll::Ready(Ok(None)),
149 None => unreachable!("unfold stream never ends"),
150 }
151 }
152}
153
154impl transport::Connection for Connection {
155 #[inline]
156 fn is_handshake_complete(&self) -> bool {
157 self.conn.handshake_data().is_some()
158 }
159
160 #[inline]
161 fn poll_shutdown(
162 &mut self,
163 _cx: &mut Context<'_>,
164 error_code: u64,
165 ) -> Poll<Result<(), TransportError>> {
166 let code = match quinn::VarInt::from_u64(error_code) {
167 Ok(code) => code,
168 Err(_) => return Poll::Ready(Err(TransportError::Other)),
169 };
170 self.conn.close(code, b"");
173 Poll::Ready(Ok(()))
174 }
175}
176
177pub(crate) struct Send {
179 stream: quinn::SendStream,
180 buf: bytes::BytesMut,
186 in_flight: Option<(usize, usize)>,
195}
196
197impl Send {
198 #[inline]
199 fn new(stream: quinn::SendStream) -> Self {
200 Self {
201 stream,
202 buf: bytes::BytesMut::new(),
203 in_flight: None,
204 }
205 }
206
207 #[inline]
208 fn id(&self) -> u64 {
209 self.stream.id().into()
210 }
211
212 #[inline]
213 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
214 let id = (data.as_ptr() as usize, data.len());
218 if !data.is_empty() && self.in_flight != Some(id) {
219 self.buf.extend_from_slice(data);
220 self.in_flight = Some(id);
221 }
222 loop {
223 if self.buf.is_empty() {
224 self.in_flight = None;
225 return Poll::Ready(Ok(()));
226 }
227 let written = ready!(Pin::new(&mut self.stream).poll_write(cx, &self.buf))
232 .map_err(map_write_error)?;
233 debug_assert!(written > 0);
234 self.buf.advance(written);
235 }
236 }
237
238 #[inline]
239 fn poll_finish(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
240 let _ = self.stream.finish();
244 Poll::Ready(Ok(()))
245 }
246
247 #[inline]
248 fn poll_reset(&mut self, _cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
249 let code = match quinn::VarInt::from_u64(code) {
250 Ok(code) => code,
251 Err(_) => return Poll::Ready(Err(TransportError::Other)),
252 };
253 if self.stream.reset(code).is_err() {
254 return Poll::Ready(Err(TransportError::Other));
255 }
256 Poll::Ready(Ok(()))
257 }
258}
259
260type ReadChunkFuture = ReusableBoxFuture<
261 'static,
262 (
263 quinn::RecvStream,
264 Result<Option<quinn::Chunk>, quinn::ReadError>,
265 ),
266>;
267
268pub(crate) struct Recv {
270 stream: Option<quinn::RecvStream>,
271 read_chunk_fut: ReadChunkFuture,
277}
278
279impl Recv {
280 #[inline]
281 fn new(stream: quinn::RecvStream) -> Self {
282 Self {
283 stream: Some(stream),
284 read_chunk_fut: ReusableBoxFuture::new(async { unreachable!("armed before poll") }),
285 }
286 }
287
288 #[inline]
289 fn id(&self) -> u64 {
290 self.stream.as_ref().map_or(0, |stream| stream.id().into())
291 }
292
293 #[inline]
294 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
295 if let Some(mut stream) = self.stream.take() {
296 self.read_chunk_fut.set(async move {
297 let chunk = stream.read_chunk(usize::MAX, true).await;
298 (stream, chunk)
299 });
300 }
301 let (stream, chunk) = ready!(self.read_chunk_fut.poll(cx));
302 self.stream = Some(stream);
303 Poll::Ready(
304 chunk
305 .map_err(map_read_error)
306 .map(|chunk| chunk.map(|chunk| chunk.bytes)),
307 )
308 }
309
310 #[inline]
311 fn stop_sending(&mut self, code: u64) -> Result<(), TransportError> {
312 let code = match quinn::VarInt::from_u64(code) {
313 Ok(code) => code,
314 Err(_) => return Err(TransportError::Other),
315 };
316 match self.stream.as_mut() {
317 Some(stream) => stream.stop(code).map_err(|_| TransportError::Other),
318 None => Err(TransportError::Other),
319 }
320 }
321}
322
323pub(crate) struct BidiStream {
325 send: Send,
326 recv: Recv,
327}
328
329impl RecvStreamTrait for BidiStream {
330 #[inline]
331 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
332 self.recv.poll_recv(cx)
333 }
334
335 #[inline]
336 fn id(&self) -> u64 {
337 self.send.id()
338 }
339}
340
341impl SendStreamTrait for BidiStream {
342 #[inline]
343 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
344 self.send.poll_send(cx, data)
345 }
346
347 #[inline]
348 fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
349 self.send.poll_finish(cx)
350 }
351
352 #[inline]
353 fn poll_reset(&mut self, cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
354 self.send.poll_reset(cx, code)
355 }
356
357 #[inline]
358 fn poll_stop_sending(
359 &mut self,
360 _cx: &mut Context<'_>,
361 code: u64,
362 ) -> Poll<Result<(), TransportError>> {
363 Poll::Ready(self.recv.stop_sending(code))
367 }
368}
369
370impl BidiStreamTrait for BidiStream {}
371
372pub(crate) enum UniStream {
381 Send(Send),
383 Recv(Recv),
385}
386
387impl RecvStreamTrait for UniStream {
388 #[inline]
389 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
390 match self {
391 UniStream::Recv(recv) => recv.poll_recv(cx),
392 UniStream::Send(_) => Poll::Ready(Err(TransportError::Transport)),
393 }
394 }
395
396 #[inline]
397 fn id(&self) -> u64 {
398 match self {
399 UniStream::Send(send) => send.id(),
400 UniStream::Recv(recv) => recv.id(),
401 }
402 }
403}
404
405impl SendStreamTrait for UniStream {
406 #[inline]
407 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
408 match self {
409 UniStream::Send(send) => send.poll_send(cx, data),
410 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
411 }
412 }
413
414 #[inline]
415 fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
416 match self {
417 UniStream::Send(send) => send.poll_finish(cx),
418 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
419 }
420 }
421
422 #[inline]
423 fn poll_reset(&mut self, cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
424 match self {
425 UniStream::Send(send) => send.poll_reset(cx, code),
426 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
427 }
428 }
429
430 #[inline]
431 fn poll_stop_sending(
432 &mut self,
433 _cx: &mut Context<'_>,
434 code: u64,
435 ) -> Poll<Result<(), TransportError>> {
436 match self {
437 UniStream::Recv(recv) => Poll::Ready(recv.stop_sending(code)),
438 UniStream::Send(_) => Poll::Ready(Err(TransportError::Transport)),
439 }
440 }
441}
442
443impl UniStreamTrait for UniStream {}
444
445#[inline]
446fn map_connection_error(err: quinn::ConnectionError) -> TransportError {
447 match err {
448 quinn::ConnectionError::ApplicationClosed(close) => TransportError::Closed {
449 code: close.error_code.into_inner(),
450 },
451 quinn::ConnectionError::TimedOut => TransportError::Timeout,
452 _ => TransportError::Transport,
455 }
456}
457
458#[inline]
459fn map_read_error(err: quinn::ReadError) -> TransportError {
460 match err {
461 quinn::ReadError::Reset(code) => TransportError::Reset {
462 code: code.into_inner(),
463 },
464 quinn::ReadError::ConnectionLost(err) => map_connection_error(err),
465 quinn::ReadError::ClosedStream => TransportError::Closed { code: 0 },
467 quinn::ReadError::IllegalOrderedRead => TransportError::Transport,
469 quinn::ReadError::ZeroRttRejected => TransportError::Transport,
470 }
471}
472
473#[inline]
474fn map_write_error(err: quinn::WriteError) -> TransportError {
475 match err {
476 quinn::WriteError::Stopped(code) => TransportError::Stopped {
477 code: code.into_inner(),
478 },
479 quinn::WriteError::ConnectionLost(err) => map_connection_error(err),
480 quinn::WriteError::ClosedStream => TransportError::Closed { code: 0 },
482 quinn::WriteError::ZeroRttRejected => TransportError::Transport,
483 }
484}