noq_proto/connection/
state.rs

1use bytes::Bytes;
2use tracing::trace;
3
4use crate::frame::Close;
5use crate::shared::EndpointEventInner;
6use crate::{ApplicationClose, ConnectionClose, ConnectionError, TransportError, TransportErrorCode};
7
8#[allow(unreachable_pub)] // fuzzing only
9#[derive(Debug, Clone)]
10pub struct State {
11    /// Nested [`InnerState`] to enforce all state transitions are done in this module.
12    inner: InnerState,
13}
14
15impl State {
16    pub(super) fn as_handshake_mut(&mut self) -> Option<&mut Handshake> {
17        if let InnerState::Handshake(ref mut hs) = self.inner {
18            Some(hs)
19        } else {
20            None
21        }
22    }
23
24    pub(super) fn as_handshake(&self) -> Option<&Handshake> {
25        if let InnerState::Handshake(ref hs) = self.inner {
26            Some(hs)
27        } else {
28            None
29        }
30    }
31
32    pub(super) fn as_closed(&self) -> Option<&CloseReason> {
33        if let InnerState::Closed {
34            ref remote_reason, ..
35        } = self.inner
36        {
37            Some(remote_reason)
38        } else {
39            None
40        }
41    }
42
43    #[allow(unreachable_pub)] // fuzzing only
44    #[cfg(any(test, fuzzing))]
45    pub fn established() -> Self {
46        Self {
47            inner: InnerState::Established,
48        }
49    }
50
51    pub(super) fn handshake(hs: Handshake) -> Self {
52        Self {
53            inner: InnerState::Handshake(hs),
54        }
55    }
56
57    pub(super) fn move_to_handshake(&mut self, hs: Handshake) {
58        self.inner = InnerState::Handshake(hs);
59        trace!("connection state: handshake");
60    }
61
62    pub(super) fn move_to_established(&mut self) {
63        self.inner = InnerState::Established;
64        trace!("connection state: established");
65    }
66
67    /// Moves to the drained state.
68    ///
69    /// Panics if the state was already drained.
70    ///
71    /// Emits the appropriate `Draining` and `Drained` endpoint events into `events`.
72    pub(super) fn move_to_drained(
73        &mut self,
74        error: Option<ConnectionError>,
75        events: &mut impl Extend<EndpointEventInner>,
76    ) {
77        let (error, is_local, was_draining) = if let Some(error) = error {
78            (
79                Some(error),
80                false,
81                matches!(self.inner, InnerState::Draining { .. }),
82            )
83        } else {
84            let (error, was_draining) = match &mut self.inner {
85                InnerState::Draining { error, .. } => (error.take(), true),
86                InnerState::Drained { .. } => panic!("invalid state transition drained -> drained"),
87                InnerState::Closed { error_read, .. } if *error_read => (None, false),
88                InnerState::Closed { remote_reason, .. } => {
89                    let error = match remote_reason.clone().into() {
90                        ConnectionError::ConnectionClosed(close) => {
91                            if close.error_code == TransportErrorCode::PROTOCOL_VIOLATION {
92                                ConnectionError::TransportError(TransportError::new(
93                                    close.error_code,
94                                    String::from_utf8_lossy(&close.reason[..]).to_string(),
95                                ))
96                            } else {
97                                ConnectionError::ConnectionClosed(close)
98                            }
99                        }
100                        e => e,
101                    };
102                    (Some(error), false)
103                }
104                InnerState::Handshake(_) | InnerState::Established => (None, false),
105            };
106            (error, self.is_local_close(), was_draining)
107        };
108        self.inner = InnerState::Drained { error, is_local };
109        trace!("connection state: drained");
110
111        if !was_draining {
112            events.extend([EndpointEventInner::Draining]);
113        }
114        events.extend([EndpointEventInner::Drained]);
115    }
116
117    /// Moves to a draining state.
118    ///
119    /// Panics if the state is already draining or drained.
120    ///
121    /// Emits a `Draining` endpoint event into `events`.
122    pub(super) fn move_to_draining(
123        &mut self,
124        error: Option<ConnectionError>,
125        events: &mut impl Extend<EndpointEventInner>,
126    ) {
127        assert!(
128            matches!(
129                self.inner,
130                InnerState::Handshake(_) | InnerState::Established | InnerState::Closed { .. }
131            ),
132            "invalid state transition {:?} -> draining",
133            self.as_type()
134        );
135        let is_local = self.is_local_close();
136
137        // If no error is provided and we were in the `Closed` state with an unread non-local
138        // error, preserve that error so `take_error` can still surface it. Otherwise the
139        // `ConnectionLost` event would never be emitted, breaking the invariant that a
140        // drained noq-proto connection will always produce a `ConnectionLost` event.
141        let error = error.or_else(|| {
142            if let InnerState::Closed {
143                ref remote_reason,
144                error_read: false,
145                is_local: false,
146            } = self.inner
147            {
148                Some(remote_reason.clone().into())
149            } else {
150                None
151            }
152        });
153
154        self.inner = InnerState::Draining { error, is_local };
155        trace!("connection state: draining");
156        events.extend([EndpointEventInner::Draining]);
157    }
158
159    fn is_local_close(&self) -> bool {
160        match self.inner {
161            InnerState::Handshake(_) => false,
162            InnerState::Established => false,
163            InnerState::Closed { is_local, .. } => is_local,
164            InnerState::Draining { is_local, .. } => is_local,
165            InnerState::Drained { is_local, .. } => is_local,
166        }
167    }
168
169    /// Enters the Closing connection state, due to changes in the [`Connection`] state.
170    ///
171    /// This is the closing state from
172    /// <https://www.rfc-editor.org/rfc/rfc9000.html#section-10.2.1> due to the local side
173    /// having initiated immediate close.
174    ///
175    /// Crucially, this is to be used when internal state changes result in initiating an
176    /// immediate close. The resulting error will be surfaced as a [`ConnectionLost`] event
177    /// in [`Connection::poll`].
178    ///
179    /// # Panics
180    ///
181    /// Panics if the state is later than established.
182    ///
183    /// [`Connection`]: super::Connection
184    /// [`ConnectionLost`]: crate::Event::ConnectionLost
185    /// [`Connection::poll`]: super::Connection::poll
186    pub(super) fn move_to_closed<R: Into<CloseReason>>(&mut self, reason: R) {
187        assert!(
188            matches!(
189                self.inner,
190                InnerState::Handshake(_) | InnerState::Established | InnerState::Closed { .. }
191            ),
192            "invalid state transition {:?} -> closed",
193            self.as_type()
194        );
195        let remote_reason = reason.into();
196        let is_local = false;
197        trace!(?remote_reason, ?is_local, "connection state: closed");
198        self.inner = InnerState::Closed {
199            error_read: false,
200            remote_reason,
201            is_local,
202        };
203    }
204
205    /// Enters the Closing connection state, initiated by explicit API calls.
206    ///
207    /// This is the closing state from
208    /// <https://www.rfc-editor.org/rfc/rfc9000.html#section-10.2.1> due to the local side
209    /// having initiated immediate close.
210    ///
211    /// Crucially, this is to be used when immediate close is entered due to an API
212    /// being called. It means the close will NOT surface as a [`ConnectionLost`] event in
213    /// [`Connection::poll`].
214    ///
215    /// See [`Self::move_to_closed`] for when the internal state changes resulted in
216    /// initiating an immediate close.
217    ///
218    /// # Panics
219    ///
220    /// Panics if the state is later than established.
221    ///
222    /// [`ConnectionLost`]: crate::Event::ConnectionLost
223    /// [`Connection::poll`]: super::Connection::poll
224    pub(super) fn move_to_closed_local<R: Into<CloseReason>>(&mut self, reason: R) {
225        assert!(
226            matches!(
227                self.inner,
228                InnerState::Handshake(_) | InnerState::Established | InnerState::Closed { .. }
229            ),
230            "invalid state transition {:?} -> closed (local)",
231            self.as_type()
232        );
233        let remote_reason = reason.into();
234        let is_local = true;
235        trace!(?remote_reason, ?is_local, "connection state: closed");
236        self.inner = InnerState::Closed {
237            error_read: false,
238            remote_reason,
239            is_local,
240        };
241    }
242
243    pub(super) fn is_handshake(&self) -> bool {
244        matches!(self.inner, InnerState::Handshake(_))
245    }
246
247    pub(super) fn is_established(&self) -> bool {
248        matches!(self.inner, InnerState::Established)
249    }
250
251    pub(super) fn is_closed(&self) -> bool {
252        matches!(
253            self.inner,
254            InnerState::Closed { .. } | InnerState::Draining { .. } | InnerState::Drained { .. }
255        )
256    }
257
258    pub(super) fn is_drained(&self) -> bool {
259        matches!(self.inner, InnerState::Drained { .. })
260    }
261
262    pub(super) fn take_error(&mut self) -> Option<ConnectionError> {
263        match &mut self.inner {
264            InnerState::Draining { error, is_local } => {
265                if !*is_local {
266                    error.take()
267                } else {
268                    None
269                }
270            }
271            InnerState::Drained { error, is_local } => {
272                if !*is_local {
273                    error.take()
274                } else {
275                    None
276                }
277            }
278            InnerState::Closed {
279                remote_reason,
280                is_local: local_reason,
281                error_read,
282            } => {
283                if *error_read {
284                    None
285                } else {
286                    *error_read = true;
287                    if *local_reason {
288                        None
289                    } else {
290                        Some(remote_reason.clone().into())
291                    }
292                }
293            }
294            InnerState::Handshake(_) | InnerState::Established => None,
295        }
296    }
297
298    pub(super) fn as_type(&self) -> StateType {
299        match self.inner {
300            InnerState::Handshake(_) => StateType::Handshake,
301            InnerState::Established => StateType::Established,
302            InnerState::Closed { .. } => StateType::Closed,
303            InnerState::Draining { .. } => StateType::Draining,
304            InnerState::Drained { .. } => StateType::Drained,
305        }
306    }
307}
308
309/// The state a [`Connection`] can be in.
310///
311/// [`Connection`]: super::Connection
312#[derive(Debug, Clone)]
313pub(super) enum StateType {
314    /// Before the handshake is *confirmed*.
315    Handshake,
316    /// Once the handshake is *confirmed*.
317    Established,
318    /// The connection is closed, waiting for remote to confirm close.
319    ///
320    /// Specifically <https://www.rfc-editor.org/rfc/rfc9000.html#section-10.2.1>.
321    ///
322    /// So the side that initiates an immediate close will stay in this state while it is
323    /// waiting for the remote to also send a CONNECTION_CLOSE. The side that receives a
324    /// connection close will skip straight to [`StateType::Draining`].
325    Closed,
326    /// The connection is draining, giving time to gracefully discard any in-flight packets.
327    ///
328    /// Specifically <https://www.rfc-editor.org/rfc/rfc9000.html#section-10.2.2>.
329    ///
330    /// See [`StateType::Closed`] above for more details.
331    Draining,
332    /// The connection is drained, waiting for the application to drop us.
333    ///
334    /// This is a terminal state in which the connection does nothing and can never do
335    /// anything again. Waiting for the application to drop the [`Connection`] struct.
336    ///
337    /// [`Connection`]: super::Connection
338    Drained,
339}
340
341#[derive(Debug, Clone)]
342pub(super) enum CloseReason {
343    TransportError(TransportError),
344    Connection(ConnectionClose),
345    Application(ApplicationClose),
346}
347
348impl From<TransportError> for CloseReason {
349    fn from(x: TransportError) -> Self {
350        Self::TransportError(x)
351    }
352}
353impl From<ConnectionClose> for CloseReason {
354    fn from(x: ConnectionClose) -> Self {
355        Self::Connection(x)
356    }
357}
358impl From<ApplicationClose> for CloseReason {
359    fn from(x: ApplicationClose) -> Self {
360        Self::Application(x)
361    }
362}
363
364impl From<Close> for CloseReason {
365    fn from(value: Close) -> Self {
366        match value {
367            Close::Application(reason) => Self::Application(reason),
368            Close::Connection(reason) => Self::Connection(reason),
369        }
370    }
371}
372
373impl From<CloseReason> for ConnectionError {
374    fn from(value: CloseReason) -> Self {
375        match value {
376            CloseReason::TransportError(err) => Self::TransportError(err),
377            CloseReason::Connection(reason) => Self::ConnectionClosed(reason),
378            CloseReason::Application(reason) => Self::ApplicationClosed(reason),
379        }
380    }
381}
382
383impl From<CloseReason> for Close {
384    fn from(value: CloseReason) -> Self {
385        match value {
386            CloseReason::TransportError(err) => Self::Connection(err.into()),
387            CloseReason::Connection(reason) => Self::Connection(reason),
388            CloseReason::Application(reason) => Self::Application(reason),
389        }
390    }
391}
392
393#[derive(Debug, Clone)]
394enum InnerState {
395    /// See [`StateType::Handshake`].
396    Handshake(Handshake),
397    /// See [`StateType::Established`].
398    Established,
399    /// See [`StateType::Closed`].
400    Closed {
401        /// The reason the remote closed the connection, or the reason we are sending to the
402        /// remote.
403        remote_reason: CloseReason,
404        /// Set to true if we closed the connection locally.
405        is_local: bool,
406        /// Did we read this as error already?
407        error_read: bool,
408    },
409    /// See [`StateType::Draining`].
410    Draining {
411        /// Why the connection was lost, if it has been.
412        error: Option<ConnectionError>,
413        /// Set to true if we closed the connection locally.
414        is_local: bool,
415    },
416    /// See [`StateType::Drained`].
417    /// Waiting for application to call close so we can dispose of the resources.
418    Drained {
419        /// Why the connection was lost, if it has been.
420        error: Option<ConnectionError>,
421        /// Set to true if we closed the connection locally.
422        is_local: bool,
423    },
424}
425
426#[allow(unreachable_pub)] // fuzzing only
427#[derive(Debug, Clone)]
428pub struct Handshake {
429    /// Whether the remote CID has been set by the peer yet.
430    ///
431    /// Always set for servers.
432    pub(super) remote_cid_set: bool,
433    /// Stateless retry token received in the first Initial by a server.
434    ///
435    /// Must be present in every Initial. Always empty for clients.
436    pub(super) expected_token: Bytes,
437    /// First cryptographic message.
438    ///
439    /// Only set for clients.
440    pub(super) client_hello: Option<Bytes>,
441    /// Whether the server address is allowed to migrate.
442    ///
443    /// We allow the server to migrate during the handshake as long as we have not
444    /// received an authenticated handshake packet: it can send a response from a
445    /// different address than we sent the initial to.  This allows us to send the
446    /// initial packet over multiple paths - by means of an IPv6 ULA address that copies
447    /// the packets sent to it to multiple destinations - and accept one response.
448    ///
449    /// This is only ever set to true if for a client which hasn't yet received an
450    /// authenticated handshake packet.  It is set back to false in
451    /// [`super::Connection::on_packet_authenticated`].
452    ///
453    /// THIS IS NOT RFC 9000 COMPLIANT!  A server is not allowed to migrate addresses,
454    /// other than using the preferred-address transport parameter.
455    pub(super) allow_server_migration: bool,
456}