iota_network_stack/
request_message_timeout.rs

1// Copyright (c) 2026 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! Deadline for a gRPC request message to arrive after its headers.
5
6use std::{
7    future::Future,
8    pin::Pin,
9    task::{Context, Poll, ready},
10    time::Duration,
11};
12
13use http::Request;
14use pin_project_lite::pin_project;
15use tokio::time::Sleep;
16use tonic::{Status, server::NamedService};
17use tower::Service;
18
19type BoxError = Box<dyn std::error::Error + Send + Sync>;
20
21/// Fails a request whose body has not fully arrived within `timeout` of its
22/// headers with `DEADLINE_EXCEEDED`, so a peer cannot keep a request open by
23/// withholding its message. Requests on `exempt_paths` pass through untouched;
24/// they carry client streams that stay open by design.
25#[derive(Debug, Clone)]
26pub struct RequestMessageTimeout<S> {
27    inner: S,
28    timeout: Duration,
29    exempt_paths: &'static [&'static str],
30}
31
32impl<S> RequestMessageTimeout<S> {
33    pub fn new(inner: S, timeout: Duration, exempt_paths: &'static [&'static str]) -> Self {
34        Self {
35            inner,
36            timeout,
37            exempt_paths,
38        }
39    }
40}
41
42impl<S, B> Service<Request<B>> for RequestMessageTimeout<S>
43where
44    S: Service<Request<TimedRequestBody<B>>>,
45{
46    type Response = S::Response;
47    type Error = S::Error;
48    type Future = S::Future;
49
50    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
51        self.inner.poll_ready(cx)
52    }
53
54    fn call(&mut self, request: Request<B>) -> Self::Future {
55        let timeout = (!self.exempt_paths.contains(&request.uri().path())).then_some(self.timeout);
56        self.inner
57            .call(request.map(|body| TimedRequestBody::new(body, timeout)))
58    }
59}
60
61impl<S: NamedService> NamedService for RequestMessageTimeout<S> {
62    const NAME: &'static str = S::NAME;
63}
64
65pin_project! {
66    /// Request body that yields `DEADLINE_EXCEEDED` in place of its next frame
67    /// once its deadline has passed.
68    pub struct TimedRequestBody<B> {
69        #[pin]
70        inner: B,
71        #[pin]
72        deadline: Option<Sleep>,
73    }
74}
75
76impl<B> TimedRequestBody<B> {
77    fn new(inner: B, timeout: Option<Duration>) -> Self {
78        Self {
79            inner,
80            deadline: timeout.map(tokio::time::sleep),
81        }
82    }
83}
84
85impl<B> http_body::Body for TimedRequestBody<B>
86where
87    B: http_body::Body,
88    B::Error: Into<BoxError>,
89{
90    type Data = B::Data;
91    type Error = BoxError;
92
93    fn poll_frame(
94        self: Pin<&mut Self>,
95        cx: &mut Context<'_>,
96    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
97        let this = self.project();
98        if let Poll::Ready(frame) = this.inner.poll_frame(cx) {
99            return Poll::Ready(frame.map(|result| result.map_err(Into::into)));
100        }
101        if let Some(deadline) = this.deadline.as_pin_mut() {
102            ready!(deadline.poll(cx));
103            return Poll::Ready(Some(Err(Box::new(Status::deadline_exceeded(
104                "request message did not arrive in time",
105            )))));
106        }
107        Poll::Pending
108    }
109
110    fn is_end_stream(&self) -> bool {
111        self.inner.is_end_stream()
112    }
113
114    fn size_hint(&self) -> http_body::SizeHint {
115        self.inner.size_hint()
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use std::{
122        convert::Infallible,
123        future::{Ready, ready},
124    };
125
126    use bytes::Bytes;
127    use http_body_util::{BodyExt, Full};
128
129    use super::*;
130
131    const TIMEOUT: Duration = Duration::from_millis(50);
132    const EXEMPT: &[&str] = &["/pkg.Svc/Stream"];
133
134    /// A body whose next frame never arrives.
135    struct Withheld;
136
137    impl http_body::Body for Withheld {
138        type Data = Bytes;
139        type Error = Infallible;
140
141        fn poll_frame(
142            self: Pin<&mut Self>,
143            _: &mut Context<'_>,
144        ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
145            Poll::Pending
146        }
147    }
148
149    /// Hands the request body back, so a test can drive it as tonic would.
150    #[derive(Clone)]
151    struct Echo;
152
153    impl<B> Service<Request<B>> for Echo {
154        type Response = B;
155        type Error = Infallible;
156        type Future = Ready<Result<B, Infallible>>;
157
158        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
159            Poll::Ready(Ok(()))
160        }
161
162        fn call(&mut self, request: Request<B>) -> Self::Future {
163            ready(Ok(request.into_body()))
164        }
165    }
166
167    fn request<B>(path: &str, body: B) -> Request<B> {
168        Request::builder().uri(path).body(body).unwrap()
169    }
170
171    #[tokio::test]
172    async fn withheld_message_fails_with_deadline_exceeded() {
173        let mut service = RequestMessageTimeout::new(Echo, TIMEOUT, EXEMPT);
174        let body = service
175            .call(request("/pkg.Svc/Unary", Withheld))
176            .await
177            .unwrap();
178        let error = body.collect().await.unwrap_err();
179        assert_eq!(
180            Status::from_error(error).code(),
181            tonic::Code::DeadlineExceeded
182        );
183    }
184
185    #[tokio::test]
186    async fn withheld_message_on_an_exempt_path_stays_open() {
187        let mut service = RequestMessageTimeout::new(Echo, TIMEOUT, EXEMPT);
188        let body = service
189            .call(request("/pkg.Svc/Stream", Withheld))
190            .await
191            .unwrap();
192        assert!(
193            tokio::time::timeout(TIMEOUT * 3, body.collect())
194                .await
195                .is_err(),
196            "an exempt request must not be cut by the deadline"
197        );
198    }
199
200    #[tokio::test]
201    async fn message_that_arrives_in_time_passes_through() {
202        let mut service = RequestMessageTimeout::new(Echo, TIMEOUT, EXEMPT);
203        let body = service
204            .call(request(
205                "/pkg.Svc/Unary",
206                Full::new(Bytes::from_static(b"message")),
207            ))
208            .await
209            .unwrap();
210        assert_eq!(body.collect().await.unwrap().to_bytes(), "message");
211    }
212}
iota_network_stack/request_message_timeout.rs

iota_network_stack/
request_message_timeout.rs
