Check h2 stream end for error

Certain APIs consuming HTTPTasks want to know if the response has
finished. However, in cases such as protocol errors the h2 RecvStream
will return `is_end_stream` which may be associated with an HTTPTask
that was not actually the end of stream. This affects systems such as
caching which wants to know when the response has properly finished for
finishing a cache miss.

This change attempts to check if the stream has ended due to an h2 error
or not to forward a successful or Failed HTTPTask as appropriate.

1 files changed, 32 insertions, 12 deletions

diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs
index 53b2c14..fad91c2 100644
--- a/pingora-proxy/src/proxy_h2.rs
+++ b/pingora-proxy/src/proxy_h2.rs
@@ -569,9 +569,21 @@ pub(crate) async fn pipe_2to1_response(
 
     let resp_header = Box::new(client.response_header().expect("just read").clone());
 
-    tx.send(HttpTask::Header(resp_header, client.response_finished()))
-        .await
-        .or_err(InternalError, "sending h2 headers to pipe")?;
+    match client.check_response_end_or_error() {
+        Ok(eos) => {
+            tx.send(HttpTask::Header(resp_header, eos))
+                .await
+                .or_err(InternalError, "sending h2 headers to pipe")?;
+        }
+        Err(e) => {
+            // If upstream errored, then push error to downstream and then quit
+            // Don't care if send fails (which means downstream already gone)
+            // we were still able to retrieve the headers, so try sending
+            let _ = tx.send(HttpTask::Header(resp_header, false)).await;
+            let _ = tx.send(HttpTask::Failed(e.into_up())).await;
+            return Ok(());
+        }
+    }
 
     while let Some(chunk) = client
         .read_response_body()
@@ -583,21 +595,29 @@ pub(crate) async fn pipe_2to1_response(
             Ok(d) => d,
             Err(e) => {
                 // Push the error to downstream and then quit
-                // Don't care if send fails: downstream already gone
                 let _ = tx.send(HttpTask::Failed(e.into_up())).await;
                 // Downstream should consume all remaining data and handle the error
                 return Ok(());
             }
         };
-        if data.is_empty() && !client.response_finished() {
-            /* it is normal to get 0 bytes because of multi-chunk
-             * don't write 0 bytes to downstream since it will be
-             * misread as the terminating chunk */
-            continue;
+        match client.check_response_end_or_error() {
+            Ok(eos) => {
+                if data.is_empty() && !eos {
+                    /* it is normal to get 0 bytes because of multi-chunk
+                     * don't write 0 bytes to downstream since it will be
+                     * misread as the terminating chunk */
+                    continue;
+                }
+                tx.send(HttpTask::Body(Some(data), eos))
+                    .await
+                    .or_err(InternalError, "sending h2 body to pipe")?;
+            }
+            Err(e) => {
+                // Similar to above, push the error to downstream and then quit
+                let _ = tx.send(HttpTask::Failed(e.into_up())).await;
+                return Ok(());
+            }
         }
-        tx.send(HttpTask::Body(Some(data), client.response_finished()))
-            .await
-            .or_err(InternalError, "sending h2 body to pipe")?;
     }
 
     // attempt to get trailers