1use std::{
30 future::Future,
31 pin::Pin,
32 task::{Context, Poll},
33};
34
35use bytes::{Buf, Bytes};
36use futures_util::{ready, stream, Stream, StreamExt};
37use tokio_util::sync::ReusableBoxFuture;
38
39use crate::h3::error::TransportError;
40use crate::h3::transport::{
41 self, Accept, BidiStream as BidiStreamTrait, OpenStreams, RecvStream as RecvStreamTrait,
42 SendStream as SendStreamTrait, UniStream as UniStreamTrait,
43};
44
45type BoxStream<T> = Pin<Box<dyn Stream<Item = T> + std::marker::Send + 'static>>;
47type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + std::marker::Send + 'static>>;
48
49type AcceptBiItem = Result<(quinn::SendStream, quinn::RecvStream), quinn::ConnectionError>;
50type AcceptUniItem = Result<quinn::RecvStream, quinn::ConnectionError>;
51type OpenUniItem = Result<quinn::SendStream, quinn::ConnectionError>;
52
53pub struct Connection {
58 conn: quinn::Connection,
59 accept_bi: BoxStream<AcceptBiItem>,
66 accept_uni: BoxStream<AcceptUniItem>,
67 open_uni: Option<BoxStream<OpenUniItem>>,
69}
70
71impl Connection {
72 #[inline]
78 pub fn new(conn: quinn::Connection) -> Self {
79 let accept_bi = Box::pin(stream::unfold(conn.clone(), |conn| async move {
80 Some((conn.accept_bi().await, conn))
81 }));
82 let accept_uni = Box::pin(stream::unfold(conn.clone(), |conn| async move {
83 Some((conn.accept_uni().await, conn))
84 }));
85 Self {
86 conn,
87 accept_bi,
88 accept_uni,
89 open_uni: None,
90 }
91 }
92
93 #[inline]
99 pub fn close_reason(&self) -> Option<quinn::ConnectionError> {
100 self.conn.close_reason()
101 }
102}
103
104impl OpenStreams for Connection {
105 #[inline]
106 fn poll_open_uni(
107 &mut self,
108 cx: &mut Context<'_>,
109 ) -> Poll<Result<Box<dyn UniStreamTrait>, TransportError>> {
110 let stream = self.open_uni.get_or_insert_with(|| {
111 let conn = self.conn.clone();
112 Box::pin(stream::unfold(conn, |conn| async move {
113 Some((conn.open_uni().await, conn))
114 }))
115 });
116 match ready!(stream.poll_next_unpin(cx)) {
117 Some(Ok(stream)) => Poll::Ready(Ok(Box::new(UniStream::Send(Send::new(stream))))),
118 Some(Err(err)) => Poll::Ready(Err(map_connection_error(err))),
119 None => unreachable!("unfold stream never ends"),
120 }
121 }
122}
123
124impl Accept for Connection {
125 #[inline]
126 fn poll_accept(
127 &mut self,
128 cx: &mut Context<'_>,
129 ) -> Poll<Result<Option<Box<dyn BidiStreamTrait>>, TransportError>> {
130 match ready!(self.accept_bi.as_mut().poll_next_unpin(cx)) {
131 Some(Ok((send, recv))) => Poll::Ready(Ok(Some(Box::new(BidiStream {
132 send: Send::new(send),
133 recv: Recv::new(recv),
134 })))),
135 Some(Err(_)) => Poll::Ready(Ok(None)),
139 None => unreachable!("unfold stream never ends"),
140 }
141 }
142
143 #[inline]
144 fn poll_accept_uni(
145 &mut self,
146 cx: &mut Context<'_>,
147 ) -> Poll<Result<Option<Box<dyn UniStreamTrait>>, TransportError>> {
148 match ready!(self.accept_uni.as_mut().poll_next_unpin(cx)) {
149 Some(Ok(stream)) => Poll::Ready(Ok(Some(Box::new(UniStream::Recv(Recv::new(stream)))))),
150 Some(Err(_)) => Poll::Ready(Ok(None)),
151 None => unreachable!("unfold stream never ends"),
152 }
153 }
154}
155
156impl transport::Connection for Connection {
157 #[inline]
158 fn is_handshake_complete(&self) -> bool {
159 self.conn.handshake_data().is_some()
160 }
161
162 #[inline]
163 fn poll_shutdown(
164 &mut self,
165 _cx: &mut Context<'_>,
166 error_code: u64,
167 ) -> Poll<Result<(), TransportError>> {
168 let code = match quinn::VarInt::from_u64(error_code) {
169 Ok(code) => code,
170 Err(_) => return Poll::Ready(Err(TransportError::Other)),
171 };
172 self.conn.close(code, b"");
175 Poll::Ready(Ok(()))
176 }
177}
178
179pub(crate) struct Send {
181 stream: quinn::SendStream,
182 buf: bytes::BytesMut,
188 in_flight: Option<(usize, usize)>,
197 stopped_future: Option<BoxFuture<Result<Option<quinn::VarInt>, quinn::StoppedError>>>,
198}
199
200impl Send {
201 #[inline]
202 fn new(stream: quinn::SendStream) -> Self {
203 Self {
204 stream,
205 buf: bytes::BytesMut::new(),
206 in_flight: None,
207 stopped_future: None,
208 }
209 }
210
211 #[inline]
212 fn id(&self) -> u64 {
213 self.stream.id().into()
214 }
215
216 #[inline]
217 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
218 let id = (data.as_ptr() as usize, data.len());
222 if !data.is_empty() && self.in_flight != Some(id) {
223 self.buf.extend_from_slice(data);
224 self.in_flight = Some(id);
225 }
226 loop {
227 if self.buf.is_empty() {
228 self.in_flight = None;
229 return Poll::Ready(Ok(()));
230 }
231 let written = ready!(Pin::new(&mut self.stream).poll_write(cx, &self.buf))
236 .map_err(map_write_error)?;
237 debug_assert!(written > 0);
238 self.buf.advance(written);
239 }
240 }
241
242 #[inline]
243 fn poll_stopped(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
244 let mut fut = self
245 .stopped_future
246 .get_or_insert_with(|| Box::pin(self.stream.stopped()));
247 match Pin::new(&mut fut).poll(cx) {
248 Poll::Pending => Poll::Pending,
249 Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
250 Poll::Ready(Err(_)) => Poll::Ready(Err(TransportError::Other)),
251 }
252 }
253
254 #[inline]
255 fn poll_finish(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
256 let _ = self.stream.finish();
260 Poll::Ready(Ok(()))
261 }
262
263 #[inline]
264 fn poll_reset(&mut self, _cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
265 let code = match quinn::VarInt::from_u64(code) {
266 Ok(code) => code,
267 Err(_) => return Poll::Ready(Err(TransportError::Other)),
268 };
269 if self.stream.reset(code).is_err() {
270 return Poll::Ready(Err(TransportError::Other));
271 }
272 Poll::Ready(Ok(()))
273 }
274}
275
276type ReadChunkFuture = ReusableBoxFuture<
277 'static,
278 (
279 quinn::RecvStream,
280 Result<Option<quinn::Chunk>, quinn::ReadError>,
281 ),
282>;
283
284pub(crate) struct Recv {
286 stream: Option<quinn::RecvStream>,
287 read_chunk_fut: ReadChunkFuture,
293}
294
295impl Recv {
296 #[inline]
297 fn new(stream: quinn::RecvStream) -> Self {
298 Self {
299 stream: Some(stream),
300 read_chunk_fut: ReusableBoxFuture::new(async { unreachable!("armed before poll") }),
301 }
302 }
303
304 #[inline]
305 fn id(&self) -> u64 {
306 self.stream.as_ref().map_or(0, |stream| stream.id().into())
307 }
308
309 #[inline]
310 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
311 if let Some(mut stream) = self.stream.take() {
312 self.read_chunk_fut.set(async move {
313 let chunk = stream.read_chunk(usize::MAX, true).await;
314 (stream, chunk)
315 });
316 }
317 let (stream, chunk) = ready!(self.read_chunk_fut.poll(cx));
318 self.stream = Some(stream);
319 Poll::Ready(
320 chunk
321 .map_err(map_read_error)
322 .map(|chunk| chunk.map(|chunk| chunk.bytes)),
323 )
324 }
325
326 #[inline]
327 fn stop_sending(&mut self, code: u64) -> Result<(), TransportError> {
328 let code = match quinn::VarInt::from_u64(code) {
329 Ok(code) => code,
330 Err(_) => return Err(TransportError::Other),
331 };
332 match self.stream.as_mut() {
333 Some(stream) => stream.stop(code).map_err(|_| TransportError::Other),
334 None => Err(TransportError::Other),
335 }
336 }
337}
338
339pub(crate) struct BidiStream {
341 send: Send,
342 recv: Recv,
343}
344
345impl RecvStreamTrait for BidiStream {
346 #[inline]
347 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
348 self.recv.poll_recv(cx)
349 }
350
351 #[inline]
352 fn id(&self) -> u64 {
353 self.send.id()
354 }
355}
356
357impl SendStreamTrait for BidiStream {
358 #[inline]
359 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
360 self.send.poll_send(cx, data)
361 }
362
363 #[inline]
364 fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
365 self.send.poll_finish(cx)
366 }
367
368 #[inline]
369 fn poll_stopped(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
370 self.send.poll_stopped(cx)
371 }
372
373 #[inline]
374 fn poll_reset(&mut self, cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
375 self.send.poll_reset(cx, code)
376 }
377
378 #[inline]
379 fn poll_stop_sending(
380 &mut self,
381 _cx: &mut Context<'_>,
382 code: u64,
383 ) -> Poll<Result<(), TransportError>> {
384 Poll::Ready(self.recv.stop_sending(code))
388 }
389}
390
391impl BidiStreamTrait for BidiStream {}
392
393pub(crate) enum UniStream {
402 Send(Send),
404 Recv(Recv),
406}
407
408impl RecvStreamTrait for UniStream {
409 #[inline]
410 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, TransportError>> {
411 match self {
412 UniStream::Recv(recv) => recv.poll_recv(cx),
413 UniStream::Send(_) => Poll::Ready(Err(TransportError::Transport)),
414 }
415 }
416
417 #[inline]
418 fn id(&self) -> u64 {
419 match self {
420 UniStream::Send(send) => send.id(),
421 UniStream::Recv(recv) => recv.id(),
422 }
423 }
424}
425
426impl SendStreamTrait for UniStream {
427 #[inline]
428 fn poll_send(&mut self, cx: &mut Context<'_>, data: &[u8]) -> Poll<Result<(), TransportError>> {
429 match self {
430 UniStream::Send(send) => send.poll_send(cx, data),
431 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
432 }
433 }
434
435 #[inline]
436 fn poll_stopped(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
437 match self {
438 UniStream::Send(send) => send.poll_stopped(cx),
439 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
440 }
441 }
442
443 #[inline]
444 fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
445 match self {
446 UniStream::Send(send) => send.poll_finish(cx),
447 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
448 }
449 }
450
451 #[inline]
452 fn poll_reset(&mut self, cx: &mut Context<'_>, code: u64) -> Poll<Result<(), TransportError>> {
453 match self {
454 UniStream::Send(send) => send.poll_reset(cx, code),
455 UniStream::Recv(_) => Poll::Ready(Err(TransportError::Transport)),
456 }
457 }
458
459 #[inline]
460 fn poll_stop_sending(
461 &mut self,
462 _cx: &mut Context<'_>,
463 code: u64,
464 ) -> Poll<Result<(), TransportError>> {
465 match self {
466 UniStream::Recv(recv) => Poll::Ready(recv.stop_sending(code)),
467 UniStream::Send(_) => Poll::Ready(Err(TransportError::Transport)),
468 }
469 }
470}
471
472impl UniStreamTrait for UniStream {}
473
474#[inline]
475fn map_connection_error(err: quinn::ConnectionError) -> TransportError {
476 match err {
477 quinn::ConnectionError::ApplicationClosed(close) => TransportError::Closed {
478 code: close.error_code.into_inner(),
479 },
480 quinn::ConnectionError::TimedOut => TransportError::Timeout,
481 _ => TransportError::Transport,
484 }
485}
486
487#[inline]
488fn map_read_error(err: quinn::ReadError) -> TransportError {
489 match err {
490 quinn::ReadError::Reset(code) => TransportError::Reset {
491 code: code.into_inner(),
492 },
493 quinn::ReadError::ConnectionLost(err) => map_connection_error(err),
494 quinn::ReadError::ClosedStream => TransportError::Closed { code: 0 },
496 quinn::ReadError::IllegalOrderedRead => TransportError::Transport,
498 quinn::ReadError::ZeroRttRejected => TransportError::Transport,
499 }
500}
501
502#[inline]
503fn map_write_error(err: quinn::WriteError) -> TransportError {
504 match err {
505 quinn::WriteError::Stopped(code) => TransportError::Stopped {
506 code: code.into_inner(),
507 },
508 quinn::WriteError::ConnectionLost(err) => map_connection_error(err),
509 quinn::WriteError::ClosedStream => TransportError::Closed { code: 0 },
511 quinn::WriteError::ZeroRttRejected => TransportError::Transport,
512 }
513}