fix(http/retry): PeekTrailersBody<B> retains first frame (#3947)

cratelyn · web-flow · commit ea6f407c57af · 2025-06-02T09:53:24.000-07:00
see linkerd/linkerd2#14050. this change fixes a logical bug with `linkerd_http_retry::peek_trailers::PeekTrailersBody::<B>::read_body(..)`. `read_body(..)` constructs a `PeekTrailersBody<B>`, by polling the inner body to see whether or not it can reach the end of the stream by only yielding to the asynchronous runtime once. in #3559, we restructured this middleware's internal modeling to reflect the `Frame<T>`-oriented signatures of the `http_body::Body` trait's 1.0 interface. unfortunately, this included a bug which could cause the first frame in a stream to be discarded if the second `Body::poll_frame()` call (_invoked via `now_or_never()`_) returns `Pending`. this could cause non-deterministic errors for users when sending traffic to HTTPRoutes and GRPCRoutes with retry annotations applied. this change rectifies this problem, ensuring that the first frame is not discarded when attempting to peek a body's trailers. to confirm that this works as expected, additional test coverage is introduced that confirms that the data and trailers of the inner body are passed through faithfully. --- * feat(http/retry): additional `PeekTrailersBody<B>` test coverage this commit introduces additional test coverage to `linker_http_retry::peek_trailers::PeekTrailersBody<B>`. this body middleware is used to facilitate transparent http retries, and allows callers to possibly inspect the trailers for a response, by polling an `http_body::Body`. this commit introduces additional unit test coverage that confirms that the data and trailers of the inner body are passed through faithfully. Signed-off-by: katelyn martin <kate@buoyant.io> * feat(http/retry): another `PeekTrailersBody<B>` test case this commit introduces some additional coverage for bodies that return `Pending` when polled a second time. Signed-off-by: katelyn martin <kate@buoyant.io> * fix(http/retry): `PeekTrailersBody<B>` retains first frame this commit fixes a logical bug with `linkerd_http_retry::peek_trailers::PeekTrailersBody::<B>::read_body(..)`. `read_body(..)` constructs a `PeekTrailersBody<B>`, by polling the inner body to see whether or not it can reach the end of the stream by only yielding to the asynchronous runtime once. in #3559, we restructured this middleware's internal modeling to reflect the `Frame<T>`-oriented signatures of the `http_body::Body` trait's 1.0 interface. unfortunately, this included a bug which could cause the first frame in a stream to be discarded if the second `Body::poll_frame()` call (_invoked via `now_or_never()`_) returns `Pending`. this could cause non-deterministic errors for users when sending traffic to HTTPRoutes and GRPCRoutes with retry annotations applied. this commit rectifies this problem, ensuring that the first frame is not discarded when attempting to peek a body's trailers. Signed-off-by: katelyn martin <kate@buoyant.io> --------- Signed-off-by: katelyn martin <kate@buoyant.io>
diff --git a/linkerd/http/retry/src/peek_trailers.rs b/linkerd/http/retry/src/peek_trailers.rs
@@ -182,7 +182,7 @@ impl<B: Body> PeekTrailersBody<B> {
                         },
                         // The body yielded an unknown kind of frame.
                         Some(Ok(None)) => Inner::Buffered {
-                            first: None,
+                            first: Some(Ok(first)),
                             second: None,
                             inner: body,
                         },
@@ -192,7 +192,7 @@ impl<B: Body> PeekTrailersBody<B> {
                     // that a second DATA frame is on the way, and we are no longer willing to
                     // await additional frames. There are no trailers to peek.
                     Inner::Buffered {
-                        first: None,
+                        first: Some(Ok(first)),
                         second: None,
                         inner: body,
                     }
@@ -338,6 +338,7 @@ mod tests {
     use bytes::Bytes;
     use http::{HeaderMap, HeaderValue};
     use http_body::Body;
+    use http_body_util::BodyExt;
     use linkerd_error::Error;
     use linkerd_mock_http_body::MockBody;
     use std::{ops::Not, task::Poll};
@@ -354,13 +355,28 @@ mod tests {
         Some(Ok(trls))
     }
 
+    async fn collect<B>(body: B) -> (Bytes, Option<HeaderMap>)
+    where
+        B: Body,
+        B::Error: std::fmt::Debug,
+    {
+        let coll = body.collect().await.expect("can collect");
+        let trls = coll.trailers().cloned();
+        let data = coll.to_bytes();
+        (data, trls)
+    }
+
     #[tokio::test]
     async fn cannot_peek_empty() {
         let (_guard, _handle) = linkerd_tracing::test::trace_init();
         let empty = MockBody::default();
         let peek = PeekTrailersBody::read_body(empty).await;
         assert!(peek.peek_trailers().is_none());
         assert!(peek.is_end_stream());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "");
+        assert!(trailers.is_none());
     }
 
     #[tokio::test]
@@ -370,6 +386,10 @@ mod tests {
         let peek = PeekTrailersBody::read_body(only_trailers).await;
         assert!(peek.peek_trailers().is_some());
         assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
     }
 
     #[tokio::test]
@@ -381,6 +401,10 @@ mod tests {
         let peek = PeekTrailersBody::read_body(body).await;
         assert!(peek.peek_trailers().is_some());
         assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "hello");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
     }
 
     #[tokio::test]
@@ -393,6 +417,10 @@ mod tests {
         let peek = PeekTrailersBody::read_body(body).await;
         assert!(peek.peek_trailers().is_none());
         assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "hello");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
     }
 
     #[tokio::test]
@@ -405,6 +433,10 @@ mod tests {
         let peek = PeekTrailersBody::read_body(body).await;
         assert!(peek.peek_trailers().is_some());
         assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "hello");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
     }
 
     #[tokio::test]
@@ -417,5 +449,27 @@ mod tests {
         let peek = PeekTrailersBody::read_body(body).await;
         assert!(peek.peek_trailers().is_none());
         assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "hellohello");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
+    }
+
+    #[tokio::test]
+    async fn cannot_peek_body_with_various_pending_polls() {
+        let (_guard, _handle) = linkerd_tracing::test::trace_init();
+        let body = MockBody::default()
+            .then_yield_data(Poll::Ready(data()))
+            .then_yield_data(Poll::Pending)
+            .then_yield_data(Poll::Ready(data()))
+            .then_yield_data(Poll::Pending)
+            .then_yield_trailer(Poll::Ready(trailers()));
+        let peek = PeekTrailersBody::read_body(body).await;
+        assert!(peek.peek_trailers().is_none());
+        assert!(peek.is_end_stream().not());
+
+        let (data, trailers) = collect(peek).await;
+        assert_eq!(data, "hellohello");
+        assert_eq!(trailers.unwrap().get("trailer").unwrap(), "shiny");
     }
 }
