ifc_lite_server/
write_timeout.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! Connection-level write-idle enforcement.
6//!
7//! A response-body timer cannot close a stalled connection: once the peer's
8//! receive window fills, Hyper stops polling that body. The timeout therefore
9//! belongs on the accepted transport. A pending socket write arms the clock;
10//! actual write progress clears it. Expiry drops the socket immediately, which
11//! makes Hyper drop its current frame and response body, including admission
12//! permits and response-specific encoders retained by that body.
13
14use axum::serve::Listener;
15use std::{
16    future::Future,
17    io::{self, IoSlice},
18    net::SocketAddr,
19    pin::Pin,
20    task::{Context, Poll},
21    time::Duration,
22};
23use tokio::{
24    io::{AsyncRead, AsyncWrite, ReadBuf},
25    net::{TcpListener, TcpStream},
26    time::{sleep, Sleep},
27};
28
29pub(crate) struct WriteTimeoutListener {
30    inner: TcpListener,
31    timeout: Option<Duration>,
32}
33
34impl WriteTimeoutListener {
35    pub(crate) fn new(inner: TcpListener, timeout: Duration) -> Self {
36        Self {
37            inner,
38            timeout: (!timeout.is_zero()).then_some(timeout),
39        }
40    }
41}
42
43impl Listener for WriteTimeoutListener {
44    type Io = WriteTimeoutIo<TcpStream>;
45    type Addr = SocketAddr;
46
47    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
48        let (stream, address) = <TcpListener as Listener>::accept(&mut self.inner).await;
49        (WriteTimeoutIo::new(stream, self.timeout), address)
50    }
51
52    fn local_addr(&self) -> io::Result<Self::Addr> {
53        self.inner.local_addr()
54    }
55}
56
57pub(crate) struct WriteTimeoutIo<T> {
58    inner: Option<T>,
59    timeout: Option<Duration>,
60    deadline: Option<Pin<Box<Sleep>>>,
61}
62
63impl<T> WriteTimeoutIo<T> {
64    fn new(inner: T, timeout: Option<Duration>) -> Self {
65        Self {
66            inner: Some(inner),
67            timeout,
68            deadline: None,
69        }
70    }
71
72    fn disconnected() -> io::Error {
73        io::Error::new(io::ErrorKind::NotConnected, "connection already closed")
74    }
75
76    fn poll_stall(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
77        let Some(timeout) = self.timeout else {
78            return Poll::Pending;
79        };
80        let deadline = self
81            .deadline
82            .get_or_insert_with(|| Box::pin(sleep(timeout)));
83        if deadline.as_mut().poll(cx).is_pending() {
84            return Poll::Pending;
85        }
86        self.inner.take();
87        Poll::Ready(Err(io::Error::new(
88            io::ErrorKind::TimedOut,
89            "socket write made no progress before IFC_STREAM_IDLE_TIMEOUT_SECS",
90        )))
91    }
92
93    fn note_progress(&mut self) {
94        self.deadline = None;
95    }
96}
97
98impl<T: AsyncRead + Unpin> AsyncRead for WriteTimeoutIo<T> {
99    fn poll_read(
100        mut self: Pin<&mut Self>,
101        cx: &mut Context<'_>,
102        buffer: &mut ReadBuf<'_>,
103    ) -> Poll<io::Result<()>> {
104        let Some(inner) = self.inner.as_mut() else {
105            return Poll::Ready(Err(Self::disconnected()));
106        };
107        Pin::new(inner).poll_read(cx, buffer)
108    }
109}
110
111impl<T: AsyncWrite + Unpin> AsyncWrite for WriteTimeoutIo<T> {
112    fn poll_write(
113        mut self: Pin<&mut Self>,
114        cx: &mut Context<'_>,
115        buffer: &[u8],
116    ) -> Poll<io::Result<usize>> {
117        let Some(inner) = self.inner.as_mut() else {
118            return Poll::Ready(Err(Self::disconnected()));
119        };
120        match Pin::new(inner).poll_write(cx, buffer) {
121            Poll::Ready(Ok(written)) => {
122                if written > 0 || buffer.is_empty() {
123                    self.note_progress();
124                }
125                Poll::Ready(Ok(written))
126            }
127            Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
128            Poll::Pending => self.poll_stall(cx),
129        }
130    }
131
132    fn poll_write_vectored(
133        mut self: Pin<&mut Self>,
134        cx: &mut Context<'_>,
135        buffers: &[IoSlice<'_>],
136    ) -> Poll<io::Result<usize>> {
137        let all_empty = buffers.iter().all(|buffer| buffer.is_empty());
138        let Some(inner) = self.inner.as_mut() else {
139            return Poll::Ready(Err(Self::disconnected()));
140        };
141        match Pin::new(inner).poll_write_vectored(cx, buffers) {
142            Poll::Ready(Ok(written)) => {
143                if written > 0 || all_empty {
144                    self.note_progress();
145                }
146                Poll::Ready(Ok(written))
147            }
148            Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
149            Poll::Pending => self.poll_stall(cx),
150        }
151    }
152
153    fn is_write_vectored(&self) -> bool {
154        self.inner
155            .as_ref()
156            .is_some_and(AsyncWrite::is_write_vectored)
157    }
158
159    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
160        let Some(inner) = self.inner.as_mut() else {
161            return Poll::Ready(Err(Self::disconnected()));
162        };
163        match Pin::new(inner).poll_flush(cx) {
164            Poll::Ready(result) => {
165                self.note_progress();
166                Poll::Ready(result)
167            }
168            Poll::Pending => self.poll_stall(cx).map_ok(|_| ()),
169        }
170    }
171
172    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
173        let Some(inner) = self.inner.as_mut() else {
174            return Poll::Ready(Ok(()));
175        };
176        match Pin::new(inner).poll_shutdown(cx) {
177            Poll::Ready(result) => {
178                self.note_progress();
179                Poll::Ready(result)
180            }
181            Poll::Pending => self.poll_stall(cx).map_ok(|_| ()),
182        }
183    }
184}
185
186#[cfg(test)]
187#[path = "write_timeout_tests.rs"]
188mod tests;
ifc_lite_server/write_timeout.rs

ifc_lite_server/
write_timeout.rs
