1use std::{
6 collections::HashSet,
7 future::poll_fn,
8 option::Option,
9 result::Result,
10 task::{ready, Context, Poll},
11};
12
13use bytes::Buf;
14use quic::RecvStream;
15use quic::StreamId;
16use tokio::sync::mpsc;
17
18use crate::{
19 connection::ConnectionInner,
20 error::{internal_error::InternalConnectionError, Code, ConnectionError},
21 frame::FrameStream,
22 proto::{
23 frame::{Frame, PayloadLen},
24 push::PushId,
25 },
26 quic::{self, SendStream as _},
27 shared_state::{ConnectionState, SharedState},
28 stream::BufRecvStream,
29};
30
31#[cfg(feature = "tracing")]
32use tracing::{instrument, trace, warn};
33
34use super::request::RequestResolver;
35
36pub struct Connection<C, B>
44where
45 C: quic::Connection<B>,
46 B: Buf,
47{
48 pub inner: ConnectionInner<C, B>,
50 pub(super) max_field_section_size: u64,
51 pub(super) ongoing_streams: HashSet<StreamId>,
53 pub(super) request_end_recv: mpsc::UnboundedReceiver<StreamId>,
55 pub(super) request_end_send: mpsc::UnboundedSender<StreamId>,
56 pub(super) sent_closing: Option<StreamId>,
58 pub(super) recv_closing: Option<PushId>,
60 pub(super) last_accepted_stream: Option<StreamId>,
62}
63
64impl<C, B> ConnectionState for Connection<C, B>
65where
66 C: quic::Connection<B>,
67 B: Buf,
68{
69 fn shared_state(&self) -> &SharedState {
70 &self.inner.shared
71 }
72}
73
74impl<C, B> Connection<C, B>
75where
76 C: quic::Connection<B>,
77 B: Buf,
78{
79 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
85 pub async fn new(conn: C) -> Result<Self, ConnectionError> {
86 super::builder::builder().build(conn).await
87 }
88}
89
90#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
91impl<C, B> Connection<C, B>
93where
94 C: quic::Connection<B>,
95 B: Buf,
96{
97 #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
98 pub fn create_resolver(&self, stream: FrameStream<C::BidiStream, B>) -> RequestResolver<C, B> {
100 self.create_resolver_internal(stream)
101 }
102
103 #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
105 pub fn poll_accept_request_stream(
106 &mut self,
107 cx: &mut Context<'_>,
108 ) -> Poll<Result<Option<C::BidiStream>, ConnectionError>> {
109 self.poll_accept_request_stream_internal(cx)
110 }
111}
112
113impl<C, B> Connection<C, B>
114where
115 C: quic::Connection<B>,
116 B: Buf,
117{
118 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
123 pub async fn accept(&mut self) -> Result<Option<RequestResolver<C, B>>, ConnectionError> {
124 let stream = match poll_fn(|cx| self.poll_accept_request_stream_internal(cx)).await? {
126 Some(s) => FrameStream::new(BufRecvStream::new(s)),
127 None => {
128 self.shutdown(0).await?;
131 return Ok(None);
132 }
133 };
134
135 let resolver = self.create_resolver_internal(stream);
136
137 self.inner.send_grease_frame = false;
139
140 Ok(Some(resolver))
141 }
142
143 fn create_resolver_internal(
144 &self,
145 stream: FrameStream<C::BidiStream, B>,
146 ) -> RequestResolver<C, B> {
147 RequestResolver {
148 frame_stream: stream,
149 request_end_send: self.request_end_send.clone(),
150 send_grease_frame: self.inner.send_grease_frame,
151 max_field_section_size: self.max_field_section_size,
152 shared: self.inner.shared.clone(),
153 }
154 }
155
156 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
160 pub async fn shutdown(&mut self, max_requests: usize) -> Result<(), ConnectionError> {
161 let max_id = self
162 .last_accepted_stream
163 .map(|id| id + max_requests)
164 .unwrap_or(StreamId::FIRST_REQUEST);
165
166 self.inner.shutdown(&mut self.sent_closing, max_id).await
167 }
168
169 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
174 fn poll_accept_request_stream_internal(
175 &mut self,
176 cx: &mut Context<'_>,
177 ) -> Poll<Result<Option<C::BidiStream>, ConnectionError>> {
178 let _ = self.poll_control(cx)?;
179 let _ = self.poll_requests_completion(cx);
180 loop {
181 let conn = self.inner.poll_accept_bi(cx)?;
182 return match conn {
183 Poll::Pending => {
184 let done = if conn.is_pending() {
185 self.recv_closing.is_some() && self.poll_requests_completion(cx).is_ready()
186 } else {
187 self.poll_requests_completion(cx).is_ready()
188 };
189
190 if done {
191 Poll::Ready(Ok(None))
192 } else {
193 Poll::Pending
196 }
197 }
198 Poll::Ready(mut s) => {
199 if let Some(max_id) = self.sent_closing {
203 if s.send_id() > max_id {
204 s.stop_sending(Code::H3_REQUEST_REJECTED.value());
205 s.reset(Code::H3_REQUEST_REJECTED.value());
206 if self.poll_requests_completion(cx).is_ready() {
207 break Poll::Ready(Ok(None));
208 }
209 continue;
210 }
211 }
212 self.last_accepted_stream = Some(s.send_id());
213 self.ongoing_streams.insert(s.send_id());
214 Poll::Ready(Ok(Some(s)))
215 }
216 };
217 }
218 }
219
220 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
221 pub(crate) fn poll_control(
222 &mut self,
223 cx: &mut Context<'_>,
224 ) -> Poll<Result<(), ConnectionError>> {
225 while (self.poll_next_control(cx)?).is_ready() {}
226 Poll::Pending
227 }
228
229 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
230 pub(crate) fn poll_next_control(
231 &mut self,
232 cx: &mut Context<'_>,
233 ) -> Poll<Result<Frame<PayloadLen>, ConnectionError>> {
234 let frame = ready!(self.inner.poll_control(cx))?;
235
236 match &frame {
237 Frame::Settings(_setting) => {
238 #[cfg(feature = "tracing")]
239 trace!("Got settings > {:?}", _setting);
240 ()
241 }
242 &Frame::Goaway(id) => self.inner.process_goaway(&mut self.recv_closing, id)?,
243 _frame @ Frame::MaxPushId(_) | _frame @ Frame::CancelPush(_) => {
244 #[cfg(feature = "tracing")]
245 warn!("Control frame ignored {:?}", _frame);
246
247 }
260
261 frame => {
266 return Poll::Ready(Err(self.inner.handle_connection_error(
267 InternalConnectionError::new(
268 Code::H3_FRAME_UNEXPECTED,
269 format!("on server control stream: {:?}", frame),
270 ),
271 )));
272 }
273 }
274 Poll::Ready(Ok(frame))
275 }
276
277 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
278 fn poll_requests_completion(&mut self, cx: &mut Context<'_>) -> Poll<()> {
279 loop {
280 match self.request_end_recv.poll_recv(cx) {
281 Poll::Ready(None) => return Poll::Ready(()),
283 Poll::Ready(Some(id)) => {
285 self.ongoing_streams.remove(&id);
286 }
287 Poll::Pending => {
288 if self.ongoing_streams.is_empty() {
289 return Poll::Ready(());
292 } else {
293 return Poll::Pending;
294 }
295 }
296 }
297 }
298 }
299}
300
301impl<C, B> Drop for Connection<C, B>
302where
303 C: quic::Connection<B>,
304 B: Buf,
305{
306 #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))]
307 fn drop(&mut self) {
308 self.inner.close_connection(
309 Code::H3_NO_ERROR,
310 "Connection was closed by the server".to_string(),
311 );
312 }
313}
314
315pub(super) struct RequestEnd {
329 pub(super) request_end: mpsc::UnboundedSender<StreamId>,
330 pub(super) stream_id: StreamId,
331}