tonic/codec/
decode.rs

1/*
2 *
3 * Copyright 2025 gRPC authors.
4 *
5 * Permission is hereby granted, free of charge, to any person obtaining a copy
6 * of this software and associated documentation files (the "Software"), to
7 * deal in the Software without restriction, including without limitation the
8 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
9 * sell copies of the Software, and to permit persons to whom the Software is
10 * furnished to do so, subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be included in
13 * all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 *
23 */
24
25use super::compression::{CompressionEncoding, CompressionSettings, decompress};
26use super::{BufferSettings, DEFAULT_MAX_RECV_MESSAGE_SIZE, DecodeBuf, Decoder, HEADER_SIZE};
27use crate::{Code, Status, body::Body, metadata::MetadataMap};
28use bytes::{Buf, BufMut, BytesMut};
29use http::{HeaderMap, StatusCode};
30use http_body::Body as HttpBody;
31use http_body_util::BodyExt;
32use std::{
33    fmt, future,
34    pin::Pin,
35    task::ready,
36    task::{Context, Poll},
37};
38use sync_wrapper::SyncWrapper;
39use tokio_stream::Stream;
40use tracing::{debug, trace};
41
42/// Streaming requests and responses.
43///
44/// This will wrap some inner [`Body`] and [`Decoder`] and provide an interface
45/// to fetch the message stream and trailing metadata
46pub struct Streaming<T> {
47    decoder: SyncWrapper<Box<dyn Decoder<Item = T, Error = Status> + Send + 'static>>,
48    inner: StreamingInner,
49}
50
51struct StreamingInner {
52    body: SyncWrapper<Body>,
53    state: State,
54    direction: Direction,
55    buf: BytesMut,
56    trailers: Option<HeaderMap>,
57    decompress_buf: BytesMut,
58    encoding: Option<CompressionEncoding>,
59    max_message_size: Option<usize>,
60}
61
62impl<T> Unpin for Streaming<T> {}
63
64#[derive(Debug, Clone)]
65enum State {
66    ReadHeader,
67    ReadBody {
68        compression: Option<CompressionEncoding>,
69        len: usize,
70    },
71    Error(Option<Status>),
72}
73
74#[derive(Debug, PartialEq, Eq)]
75enum Direction {
76    Request,
77    Response(StatusCode),
78    EmptyResponse,
79}
80
81impl<T> Streaming<T> {
82    /// Create a new streaming response in the grpc response format for decoding a response [Body]
83    /// into message of type T
84    pub fn new_response<B, D>(
85        decoder: D,
86        body: B,
87        status_code: StatusCode,
88        encoding: Option<CompressionEncoding>,
89        max_message_size: Option<usize>,
90    ) -> Self
91    where
92        B: HttpBody + Send + 'static,
93        B::Error: Into<crate::BoxError>,
94        D: Decoder<Item = T, Error = Status> + Send + 'static,
95    {
96        Self::new(
97            decoder,
98            body,
99            Direction::Response(status_code),
100            encoding,
101            max_message_size,
102        )
103    }
104
105    /// Create empty response. For creating responses that have no content (headers + trailers only)
106    pub fn new_empty<B, D>(decoder: D, body: B) -> Self
107    where
108        B: HttpBody + Send + 'static,
109        B::Error: Into<crate::BoxError>,
110        D: Decoder<Item = T, Error = Status> + Send + 'static,
111    {
112        Self::new(decoder, body, Direction::EmptyResponse, None, None)
113    }
114
115    /// Create a new streaming request in the grpc response format for decoding a request [Body]
116    /// into message of type T
117    pub fn new_request<B, D>(
118        decoder: D,
119        body: B,
120        encoding: Option<CompressionEncoding>,
121        max_message_size: Option<usize>,
122    ) -> Self
123    where
124        B: HttpBody + Send + 'static,
125        B::Error: Into<crate::BoxError>,
126        D: Decoder<Item = T, Error = Status> + Send + 'static,
127    {
128        Self::new(
129            decoder,
130            body,
131            Direction::Request,
132            encoding,
133            max_message_size,
134        )
135    }
136
137    fn new<B, D>(
138        decoder: D,
139        body: B,
140        direction: Direction,
141        encoding: Option<CompressionEncoding>,
142        max_message_size: Option<usize>,
143    ) -> Self
144    where
145        B: HttpBody + Send + 'static,
146        B::Error: Into<crate::BoxError>,
147        D: Decoder<Item = T, Error = Status> + Send + 'static,
148    {
149        let buffer_size = decoder.buffer_settings().buffer_size;
150        Self {
151            decoder: SyncWrapper::new(Box::new(decoder)),
152            inner: StreamingInner {
153                body: SyncWrapper::new(Body::new(
154                    body.map_frame(|frame| {
155                        frame.map_data(|mut buf| buf.copy_to_bytes(buf.remaining()))
156                    })
157                    .map_err(|err| Status::map_error(err.into())),
158                )),
159                state: State::ReadHeader,
160                direction,
161                buf: BytesMut::with_capacity(buffer_size),
162                trailers: None,
163                decompress_buf: BytesMut::new(),
164                encoding,
165                max_message_size,
166            },
167        }
168    }
169}
170
171impl StreamingInner {
172    fn decode_chunk(
173        &mut self,
174        buffer_settings: BufferSettings,
175    ) -> Result<Option<DecodeBuf<'_>>, Status> {
176        if let State::ReadHeader = self.state {
177            if self.buf.remaining() < HEADER_SIZE {
178                return Ok(None);
179            }
180
181            let compression_encoding = match self.buf.get_u8() {
182                0 => None,
183                1 => {
184                    {
185                        if self.encoding.is_some() {
186                            self.encoding
187                        } else {
188                            // https://grpc.github.io/grpc/core/md_doc_compression.html
189                            // An ill-constructed message with its Compressed-Flag bit set but lacking a grpc-encoding
190                            // entry different from identity in its metadata MUST fail with INTERNAL status,
191                            // its associated description indicating the invalid Compressed-Flag condition.
192                            return Err(Status::internal(
193                                "protocol error: received message with compressed-flag but no grpc-encoding was specified",
194                            ));
195                        }
196                    }
197                }
198                f => {
199                    trace!("unexpected compression flag");
200                    let message = if let Direction::Response(status) = self.direction {
201                        format!(
202                            "protocol error: received message with invalid compression flag: {f} (valid flags are 0 and 1) while receiving response with status: {status}"
203                        )
204                    } else {
205                        format!(
206                            "protocol error: received message with invalid compression flag: {f} (valid flags are 0 and 1), while sending request"
207                        )
208                    };
209                    return Err(Status::internal(message));
210                }
211            };
212
213            let len = self.buf.get_u32() as usize;
214            let limit = self
215                .max_message_size
216                .unwrap_or(DEFAULT_MAX_RECV_MESSAGE_SIZE);
217            if len > limit {
218                return Err(Status::out_of_range(format!(
219                    "Error, decoded message length too large: found {len} bytes, the limit is: {limit} bytes"
220                )));
221            }
222
223            self.buf.reserve(len);
224
225            self.state = State::ReadBody {
226                compression: compression_encoding,
227                len,
228            }
229        }
230
231        if let State::ReadBody { len, compression } = self.state {
232            // if we haven't read enough of the message then return and keep
233            // reading
234            if self.buf.remaining() < len || self.buf.len() < len {
235                return Ok(None);
236            }
237
238            let decode_buf = if let Some(encoding) = compression {
239                self.decompress_buf.clear();
240                let limit = self
241                    .max_message_size
242                    .unwrap_or(DEFAULT_MAX_RECV_MESSAGE_SIZE);
243                let limited_out_buf = (&mut self.decompress_buf).limit(limit);
244
245                if let Err(err) = decompress(
246                    CompressionSettings {
247                        encoding,
248                        buffer_growth_interval: buffer_settings.buffer_size,
249                    },
250                    &mut self.buf,
251                    limited_out_buf,
252                    len,
253                ) {
254                    if matches!(err.kind(), std::io::ErrorKind::WriteZero) {
255                        return Err(Status::resource_exhausted(format!(
256                            "Error decompressing: size limit, of {limit} bytes, exceeded while decompressing message"
257                        )));
258                    }
259                    let message = if let Direction::Response(status) = self.direction {
260                        format!(
261                            "Error decompressing: {err}, while receiving response with status: {status}"
262                        )
263                    } else {
264                        format!("Error decompressing: {err}, while sending request")
265                    };
266                    return Err(Status::internal(message));
267                }
268                let decompressed_len = self.decompress_buf.len();
269                DecodeBuf::new(&mut self.decompress_buf, decompressed_len)
270            } else {
271                DecodeBuf::new(&mut self.buf, len)
272            };
273
274            return Ok(Some(decode_buf));
275        }
276
277        Ok(None)
278    }
279
280    // Returns Some(()) if data was found or None if the loop in `poll_next` should break
281    fn poll_frame(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<()>, Status>> {
282        let frame = match ready!(Pin::new(self.body.get_mut()).poll_frame(cx)) {
283            Some(Ok(frame)) => frame,
284            Some(Err(status)) => {
285                if self.direction == Direction::Request && status.code() == Code::Cancelled {
286                    return Poll::Ready(Ok(None));
287                }
288
289                let _ = std::mem::replace(&mut self.state, State::Error(Some(status.clone())));
290                debug!("decoder inner stream error: {:?}", status);
291                return Poll::Ready(Err(status));
292            }
293            None => {
294                // FIXME: improve buf usage.
295                return Poll::Ready(if self.buf.has_remaining() {
296                    trace!("unexpected EOF decoding stream, state: {:?}", self.state);
297                    Err(Status::internal("Unexpected EOF decoding stream."))
298                } else {
299                    Ok(None)
300                });
301            }
302        };
303
304        Poll::Ready(if frame.is_data() {
305            self.buf.put(frame.into_data().unwrap());
306            Ok(Some(()))
307        } else if frame.is_trailers() {
308            if let Some(trailers) = &mut self.trailers {
309                trailers.extend(frame.into_trailers().unwrap());
310            } else {
311                self.trailers = Some(frame.into_trailers().unwrap());
312            }
313
314            Ok(None)
315        } else {
316            panic!("unexpected frame: {frame:?}");
317        })
318    }
319
320    fn response(&mut self) -> Result<(), Status> {
321        if let Direction::Response(status) = self.direction
322            && let Err(Some(e)) = crate::status::infer_grpc_status(self.trailers.as_ref(), status)
323        {
324            // If the trailers contain a grpc-status, then we should return that as the error
325            // and otherwise stop the stream (by taking the error state)
326            self.trailers.take();
327            return Err(e);
328        }
329        Ok(())
330    }
331}
332
333impl<T> Streaming<T> {
334    /// Fetch the next message from this stream.
335    ///
336    /// # Return value
337    ///
338    /// - `Result::Err(val)` means a gRPC error was sent by the sender instead
339    ///   of a valid response message. Refer to [`Status::code`] and
340    ///   [`Status::message`] to examine possible error causes.
341    ///
342    /// - `Result::Ok(None)` means the stream was closed by the sender and no
343    ///   more messages will be delivered. Further attempts to call
344    ///   [`Streaming::message`] will result in the same return value.
345    ///
346    /// - `Result::Ok(Some(val))` means the sender streamed a valid response
347    ///   message `val`.
348    ///
349    /// ```rust
350    /// # use tonic::{Streaming, Status, codec::Decoder};
351    /// # use std::fmt::Debug;
352    /// # async fn next_message_ex<T, D>(mut request: Streaming<T>) -> Result<(), Status>
353    /// # where T: Debug,
354    /// # D: Decoder<Item = T, Error = Status> + Send  + 'static,
355    /// # {
356    /// if let Some(next_message) = request.message().await? {
357    ///     println!("{:?}", next_message);
358    /// }
359    /// # Ok(())
360    /// # }
361    /// ```
362    pub async fn message(&mut self) -> Result<Option<T>, Status> {
363        match future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await {
364            Some(Ok(m)) => Ok(Some(m)),
365            Some(Err(e)) => Err(e),
366            None => Ok(None),
367        }
368    }
369
370    /// Fetch the trailing metadata.
371    ///
372    /// This will drain the stream of all its messages to receive the trailing
373    /// metadata. If [`Streaming::message`] returns `None` then this function
374    /// will not need to poll for trailers since the body was totally consumed.
375    ///
376    /// ```rust
377    /// # use tonic::{Streaming, Status};
378    /// # async fn trailers_ex<T>(mut request: Streaming<T>) -> Result<(), Status> {
379    /// if let Some(metadata) = request.trailers().await? {
380    ///     println!("{:?}", metadata);
381    /// }
382    /// # Ok(())
383    /// # }
384    /// ```
385    pub async fn trailers(&mut self) -> Result<Option<MetadataMap>, Status> {
386        // Shortcut to see if we already pulled the trailers in the stream step
387        // we need to do that so that the stream can error on trailing grpc-status
388        if let Some(trailers) = self.inner.trailers.take() {
389            return Ok(Some(MetadataMap::from_headers(trailers)));
390        }
391
392        // To fetch the trailers we must clear the body and drop it.
393        while self.message().await?.is_some() {}
394
395        // Since we call poll_trailers internally on poll_next we need to
396        // check if it got cached again.
397        if let Some(trailers) = self.inner.trailers.take() {
398            return Ok(Some(MetadataMap::from_headers(trailers)));
399        }
400
401        // We've polled through all the frames, and still no trailers, return None
402        Ok(None)
403    }
404
405    fn decode_chunk(&mut self) -> Result<Option<T>, Status> {
406        match self
407            .inner
408            .decode_chunk(self.decoder.get_mut().buffer_settings())?
409        {
410            Some(mut decode_buf) => match self.decoder.get_mut().decode(&mut decode_buf)? {
411                Some(msg) => {
412                    self.inner.state = State::ReadHeader;
413                    Ok(Some(msg))
414                }
415                None => Ok(None),
416            },
417            None => Ok(None),
418        }
419    }
420}
421
422impl<T> Stream for Streaming<T> {
423    type Item = Result<T, Status>;
424
425    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
426        loop {
427            // When the stream encounters an error yield that error once and then on subsequent
428            // calls to poll_next return Poll::Ready(None) indicating that the stream has been
429            // fully exhausted.
430            if let State::Error(status) = &mut self.inner.state {
431                return Poll::Ready(status.take().map(Err));
432            }
433
434            if let Some(item) = self.decode_chunk()? {
435                return Poll::Ready(Some(Ok(item)));
436            }
437
438            if ready!(self.inner.poll_frame(cx))?.is_none() {
439                match self.inner.response() {
440                    Ok(()) => return Poll::Ready(None),
441                    Err(err) => self.inner.state = State::Error(Some(err)),
442                }
443            }
444        }
445    }
446}
447
448impl<T> fmt::Debug for Streaming<T> {
449    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450        f.debug_struct("Streaming").finish()
451    }
452}
453
454#[cfg(test)]
455static_assertions::assert_impl_all!(Streaming<()>: Send, Sync);
tonic/codec/decode.rs

tonic/codec/
decode.rs
