leyline/core/
body_stream.rs1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll, ready};
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures_util::Stream;
8use tokio::sync::mpsc;
9use tokio::time::Sleep;
10use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
11
12use crate::core::config::HostPass;
13
14use crate::core::session::decompress::Decoder;
15use crate::trace::BodyWatch;
16
17pub struct BodyStream {
18 rx: mpsc::Receiver<std::io::Result<Bytes>>,
19 read_timeout: Option<Duration>,
20 idle: Option<Pin<Box<Sleep>>>,
21 body_limit: Option<Duration>,
22 body_deadline: Option<Pin<Box<Sleep>>>,
23 decoded: Option<Box<Decoded>>,
24 watch: Option<BodyWatch>,
25 shutdown: Option<Pin<Box<WaitForCancellationFutureOwned>>>,
26 shut: bool,
27 ended: bool,
28 pass: Option<HostPass>,
29}
30
31struct Decoded {
32 raw: BodyStream,
33 decoder: Decoder,
34 done: bool,
35}
36
37impl Decoded {
38 fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Option<std::io::Result<Bytes>>> {
39 loop {
40 if self.done {
41 return Poll::Ready(None);
42 }
43 let mut out = Vec::new();
44 let step = match Pin::new(&mut self.raw).poll_next(cx) {
45 Poll::Pending => return Poll::Pending,
46 Poll::Ready(Some(Err(e))) => {
47 self.done = true;
48 return Poll::Ready(Some(Err(e)));
49 }
50 Poll::Ready(Some(Ok(chunk))) => self.decoder.feed(&chunk, &mut out),
51 Poll::Ready(None) => {
52 self.done = true;
53 self.decoder.finish(&mut out)
54 }
55 };
56 if let Err(e) = step {
57 self.done = true;
58 return Poll::Ready(Some(Err(e.into_io())));
59 }
60 if !out.is_empty() {
61 return Poll::Ready(Some(Ok(Bytes::from(out))));
62 }
63 }
64 }
65}
66
67impl BodyStream {
68 pub(crate) fn new(rx: mpsc::Receiver<std::io::Result<Bytes>>) -> Self {
69 Self {
70 rx,
71 read_timeout: None,
72 idle: None,
73 body_limit: None,
74 body_deadline: None,
75 decoded: None,
76 watch: None,
77 shutdown: None,
78 shut: false,
79 ended: false,
80 pass: None,
81 }
82 }
83
84 pub(crate) fn decoded(mut raw: BodyStream, decoder: Decoder) -> Self {
85 let (_, rx) = mpsc::channel(1);
86 let watch = raw.watch.take();
87 Self {
88 watch,
89 decoded: Some(Box::new(Decoded {
90 raw,
91 decoder,
92 done: false,
93 })),
94 ..Self::new(rx)
95 }
96 }
97
98 pub(crate) fn from_bytes(buf: Bytes) -> Self {
99 let (tx, rx) = mpsc::channel(1);
100 if !buf.is_empty() {
101 tx.try_send(Ok(buf))
102 .expect("fresh channel with capacity 1 has room");
103 }
104 drop(tx);
105 Self::new(rx)
106 }
107
108 pub(crate) fn set_read_timeout(&mut self, timeout: Option<Duration>) {
109 self.read_timeout = timeout;
110 }
111
112 pub(crate) fn set_body_timeout(&mut self, timeout: Option<Duration>) {
113 self.body_limit = timeout;
114 self.body_deadline = None;
115 }
116
117 fn poll_body_deadline(&mut self, cx: &mut Context<'_>) -> Poll<std::io::Error> {
118 if let Some(limit) = self.body_limit.take() {
119 self.body_deadline = Some(Box::pin(tokio::time::sleep(limit)));
120 }
121 let Some(sleep) = self.body_deadline.as_mut() else {
122 return Poll::Pending;
123 };
124 ready!(sleep.as_mut().poll(cx));
125 self.body_deadline = None;
126 self.ended = true;
127 self.pass = None;
128 Poll::Ready(timed_out(
129 "body timeout: response body not complete within the configured body timeout",
130 ))
131 }
132
133 pub(crate) fn stop_on(&mut self, token: &CancellationToken) {
134 self.shutdown = Some(Box::pin(token.clone().cancelled_owned()));
135 }
136
137 pub(crate) fn hold(&mut self, pass: Option<HostPass>) {
138 self.pass = pass;
139 }
140
141 fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> bool {
142 if !self.shut
143 && let Some(wait) = self.shutdown.as_mut()
144 && wait.as_mut().poll(cx).is_ready()
145 {
146 self.shutdown = None;
147 self.shut = true;
148 self.pass = None;
149 }
150 self.shut
151 }
152
153 pub(crate) fn watch_end(&mut self) {
154 self.watch = BodyWatch::begin();
155 }
156
157 fn poll_raw(&mut self, cx: &mut Context<'_>) -> Poll<Option<std::io::Result<Bytes>>> {
158 match self.rx.poll_recv(cx) {
159 Poll::Ready(item) => {
160 self.idle = None;
161 Poll::Ready(item)
162 }
163 Poll::Pending => {
164 let Some(timeout) = self.read_timeout else {
165 return Poll::Pending;
166 };
167 let sleep = self
168 .idle
169 .get_or_insert_with(|| Box::pin(tokio::time::sleep(timeout)));
170 match sleep.as_mut().poll(cx) {
171 Poll::Ready(()) => {
172 self.idle = None;
173 Poll::Ready(Some(Err(timed_out(
174 "read timeout: no response body chunk within the configured read_timeout",
175 ))))
176 }
177 Poll::Pending => Poll::Pending,
178 }
179 }
180 }
181 }
182}
183
184fn timed_out(message: &'static str) -> std::io::Error {
185 std::io::Error::new(std::io::ErrorKind::TimedOut, message)
186}
187
188impl Stream for BodyStream {
189 type Item = std::io::Result<Bytes>;
190
191 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
192 let this = self.get_mut();
193 if this.ended {
194 return Poll::Ready(None);
195 }
196 if this.poll_shutdown(cx) {
197 this.ended = true;
198 return Poll::Ready(Some(Err(std::io::Error::other(
199 crate::core::error::Error::shut_down(),
200 ))));
201 }
202 if let Poll::Ready(err) = this.poll_body_deadline(cx) {
203 return Poll::Ready(Some(Err(err)));
204 }
205 let polled = match this.decoded.as_mut() {
206 Some(decoded) => decoded.poll(cx),
207 None => this.poll_raw(cx),
208 };
209 if let (Poll::Ready(item), Some(watch)) = (&polled, this.watch.as_mut()) {
210 watch.observe(item.as_ref());
211 }
212 if matches!(polled, Poll::Ready(None | Some(Err(_)))) {
213 this.pass = None;
214 this.ended = true;
215 }
216 polled
217 }
218}
219
220impl std::fmt::Debug for BodyStream {
221 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222 f.debug_struct("BodyStream").finish_non_exhaustive()
223 }
224}