tokio_quiche/quic/router/
connector.rs1use std::io;
28use std::mem;
29use std::sync::Arc;
30use std::task::Context;
31use std::task::Poll;
32use std::time::Instant;
33
34use datagram_socket::DatagramSocketSend;
35use datagram_socket::DatagramSocketSendExt;
36use datagram_socket::MaybeConnectedSocket;
37use datagram_socket::MAX_DATAGRAM_SIZE;
38use foundations::telemetry::log;
39use quiche::ConnectionId;
40use quiche::Header;
41use tokio_util::time::delay_queue::Key;
42use tokio_util::time::DelayQueue;
43
44use crate::quic::router::InitialPacketHandler;
45use crate::quic::router::NewConnection;
46use crate::quic::Incoming;
47use crate::quic::QuicheConnection;
48
49pub(crate) struct ClientConnector<Tx> {
53 socket_tx: MaybeConnectedSocket<Arc<Tx>>,
54 connection: ConnectionState,
55 timeout_queue: DelayQueue<ConnectionId<'static>>,
56}
57
58enum ConnectionState {
60 Queued(Box<QuicheConnection>),
62 Pending(PendingConnection),
64 Returned,
67}
68
69impl ConnectionState {
70 fn take_if_queued(&mut self) -> Option<Box<QuicheConnection>> {
71 match mem::replace(self, Self::Returned) {
72 Self::Queued(conn) => Some(conn),
73 state => {
74 *self = state;
75 None
76 },
77 }
78 }
79
80 fn take_if_pending_and_id_matches(
81 &mut self, scid: &ConnectionId<'static>,
82 ) -> Option<PendingConnection> {
83 match mem::replace(self, Self::Returned) {
84 Self::Pending(pending) if *scid == pending.conn.source_id() =>
85 Some(pending),
86 state => {
87 *self = state;
88 None
89 },
90 }
91 }
92}
93
94struct PendingConnection {
97 conn: Box<QuicheConnection>,
98 timeout_key: Option<Key>,
99 handshake_start_time: Instant,
100}
101
102impl<Tx> ClientConnector<Tx>
103where
104 Tx: DatagramSocketSend + Send + 'static,
105{
106 pub(crate) fn new(
107 socket_tx: Arc<Tx>, connection: Box<QuicheConnection>,
108 ) -> Self {
109 Self {
110 socket_tx: MaybeConnectedSocket::new(socket_tx),
111 connection: ConnectionState::Queued(connection),
112 timeout_queue: Default::default(),
113 }
114 }
115
116 fn set_connection_to_pending(
120 &mut self, mut conn: Box<QuicheConnection>,
121 ) -> io::Result<()> {
122 simple_conn_send(&self.socket_tx, &mut conn)?;
123
124 let timeout_key = conn.timeout_instant().map(|instant| {
125 self.timeout_queue
126 .insert_at(conn.source_id().into_owned(), instant.into())
127 });
128
129 self.connection = ConnectionState::Pending(PendingConnection {
130 conn,
131 timeout_key,
132 handshake_start_time: Instant::now(),
133 });
134
135 Ok(())
136 }
137
138 fn on_incoming(
143 &mut self, mut incoming: Incoming, hdr: Header<'static>,
144 ) -> io::Result<Option<NewConnection>> {
145 let Some(PendingConnection {
146 mut conn,
147 timeout_key,
148 handshake_start_time,
149 }) = self.connection.take_if_pending_and_id_matches(&hdr.dcid)
150 else {
151 log::debug!("Received Initial packet for unknown connection ID"; "scid" => ?hdr.dcid);
152 return Ok(None);
153 };
154
155 let recv_info = quiche::RecvInfo {
156 from: incoming.peer_addr,
157 to: incoming.local_addr,
158 };
159
160 if let Some(gro) = incoming.gro {
161 for dgram in incoming.buf.chunks_mut(gro as usize) {
162 let _ = conn.recv(dgram, recv_info);
164 }
165 } else {
166 let _ = conn.recv(&mut incoming.buf, recv_info);
168 }
169
170 if let Some(key) = timeout_key {
173 self.timeout_queue.remove(&key);
174 }
175
176 let scid = conn.source_id();
177 if conn.is_established() {
178 log::debug!("QUIC connection established"; "scid" => ?scid);
179
180 Ok(Some(NewConnection {
181 conn,
182 pending_cid: None,
183 initial_pkt: None,
184 cid_generator: None,
185 handshake_start_time,
186 }))
187 } else if conn.is_closed() {
188 let scid = conn.source_id();
189 log::error!("QUIC connection closed on_incoming"; "scid" => ?scid);
190
191 Err(io::Error::new(
192 io::ErrorKind::TimedOut,
193 format!("connection {scid:?} timed out"),
194 ))
195 } else {
196 self.set_connection_to_pending(conn).map(|()| None)
197 }
198 }
199
200 fn on_timeout(&mut self, scid: ConnectionId<'static>) -> io::Result<()> {
204 log::debug!("connection timedout"; "scid" => ?scid);
205
206 let Some(mut pending) =
207 self.connection.take_if_pending_and_id_matches(&scid)
208 else {
209 log::debug!("timedout connection missing from pending map"; "scid" => ?scid);
210 return Ok(());
211 };
212
213 pending.conn.on_timeout();
214
215 if pending.conn.is_closed() {
216 log::error!("pending connection closed on_timeout"; "scid" => ?scid);
217
218 return Err(io::Error::new(
219 io::ErrorKind::TimedOut,
220 format!("connection {scid:?} timed out"),
221 ));
222 }
223
224 self.set_connection_to_pending(pending.conn)
225 }
226
227 fn update(&mut self, cx: &mut Context) -> io::Result<()> {
230 while let Poll::Ready(Some(expired)) = self.timeout_queue.poll_expired(cx)
231 {
232 let scid = expired.into_inner();
233 self.on_timeout(scid)?;
234 }
235
236 if let Some(conn) = self.connection.take_if_queued() {
237 self.set_connection_to_pending(conn)?;
238 }
239
240 Ok(())
241 }
242}
243
244impl<Tx> InitialPacketHandler for ClientConnector<Tx>
245where
246 Tx: DatagramSocketSend + Send + 'static,
247{
248 fn update(&mut self, ctx: &mut Context<'_>) -> io::Result<()> {
249 ClientConnector::update(self, ctx)
250 }
251
252 fn handle_initials(
253 &mut self, incoming: Incoming, hdr: Header<'static>,
254 _: &mut quiche::Config,
255 ) -> io::Result<Option<NewConnection>> {
256 self.on_incoming(incoming, hdr)
257 }
258}
259
260fn simple_conn_send<Tx: DatagramSocketSend + Send + Sync + 'static>(
266 socket_tx: &MaybeConnectedSocket<Arc<Tx>>, conn: &mut QuicheConnection,
267) -> io::Result<()> {
268 let scid = conn.source_id().into_owned();
269 log::debug!("sending client Initials to peer"; "scid" => ?scid);
270
271 loop {
272 let scid = scid.clone();
273 let mut buf = [0; MAX_DATAGRAM_SIZE];
274 let send_res = conn.send(&mut buf);
275
276 let socket_clone = socket_tx.clone();
277 match send_res {
278 Ok((n, send_info)) => {
279 tokio::spawn({
280 let buf = buf[0..n].to_vec();
281 async move {
282 socket_clone.send_to(&buf, send_info.to).await.inspect_err(|error| {
283 log::error!("error sending client Initial packets to peer"; "scid" => ?scid, "peer_addr" => send_info.to, "error" => error.to_string());
284 })
285 }
286 });
287 },
288 Err(quiche::Error::Done) => break Ok(()),
289 Err(error) => {
290 log::error!("error writing packets to quiche's internal buffer"; "scid" => ?scid, "error" => error.to_string());
291 break Err(std::io::Error::other(error));
292 },
293 }
294 }
295}