From f049ca8f21180f7a7e5c535a1e5d3872f11e13d2 Mon Sep 17 00:00:00 2001
From: katelyn martin <kate@buoyant.io>
Date: Mon, 13 Jan 2025 00:00:00 +0000
Subject: [PATCH] feat(http/retry): `is_end_stream()` is true for empty bodies

this commit fixes a small, subtle bug in `PeekTrailersBody<B>`.

if wrapping an empty body, the peek body will inspect the wrong
`Option<T>` in the trailers field with the following type:

```rust
    /// The inner body's trailers, if it was terminated by a `TRAILERS` frame
    /// after 0-1 DATA frames, or an error if polling for trailers failed.
    ///
    /// Yes, this is a bit of a complex type, so let's break it down:
    /// - the outer `Option` indicates whether any trailers were received by
    ///   `WithTrailers`; if it's `None`, then we don't *know* if the response
    ///   had trailers, as it is not yet complete.
    /// - the inner `Result` and `Option` are the `Result` and `Option` returned
    ///   by `HttpBody::trailers` on the inner body. If this is `Ok(None)`, then
    ///   the body has terminated without trailers --- it is *known* to not have
    ///   trailers.
    trailers: Option<Result<Option<http::HeaderMap>, B::Error>>,
```

for an empty body, we *know* that there are no trailers, which means
that we have `Some(Ok(None))`.

consider also, the documentation of `is_end_stream()`:

> An end of stream means that both poll_data and poll_trailers will
> return None.
>
> A return value of false does not guarantee that a value will be
> returned from poll_stream or poll_trailers.

we can guarantee in this case that `poll_trailers()` will return
`Ok(None)` since we've already called it and proven that to be the case.
we *are* holding that value, after all.

this change will not affect any behavior w.r.t. what the peek body
yields, but it will mean that it reports `is_end_stream()` correctly
when it wraps an empty body.

Signed-off-by: katelyn martin <kate@buoyant.io>
---
 linkerd/http/retry/src/peek_trailers.rs | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/linkerd/http/retry/src/peek_trailers.rs b/linkerd/http/retry/src/peek_trailers.rs
index fd1a8a3589..267dcf8946 100644
--- a/linkerd/http/retry/src/peek_trailers.rs
+++ b/linkerd/http/retry/src/peek_trailers.rs
@@ -191,7 +191,18 @@ where
 
     #[inline]
     fn is_end_stream(&self) -> bool {
-        self.first_data.is_none() && self.trailers.is_none() && self.inner.is_end_stream()
+        let Self {
+            inner,
+            first_data,
+            trailers,
+        } = self;
+
+        let trailers_finished = match trailers {
+            Some(Ok(Some(_)) | Err(_)) => false,
+            None | Some(Ok(None)) => true,
+        };
+
+        first_data.is_none() && trailers_finished && inner.is_end_stream()
     }
 
     #[inline]
@@ -254,8 +265,7 @@ mod tests {
         let empty = MockBody::default();
         let peek = PeekTrailersBody::read_body(empty).await;
         assert!(peek.peek_trailers().is_none());
-        // TODO(kate): this will not return `true`.
-        // assert!(peek.is_end_stream());
+        assert!(peek.is_end_stream());
     }
 
     #[tokio::test]