Add ability to ignore informational responses when proxying downstream

5 files changed, 145 insertions, 19 deletions

diff --git a/.bleep b/.bleep
index 5765e75..a6c2bfb 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-9b88d76089e0f81c67cb502422148d4d26d4977e
\ No newline at end of file
+d112ca7464812ac6e575f46ec5d37b5da365dd82
\ No newline at end of file
diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs
index c6479e7..7e8d7e3 100644
--- a/pingora-core/src/protocols/http/server.rs
+++ b/pingora-core/src/protocols/http/server.rs
@@ -218,6 +218,19 @@ impl Session {
         }
     }
 
+    /// Sets whether we ignore writing informational responses downstream.
+    ///
+    /// For HTTP/1.1 this is a noop if the response is Upgrade or Continue and
+    /// Expect: 100-continue was set on the request.
+    ///
+    /// This is a noop for h2 because informational responses are always ignored.
+    pub fn set_ignore_info_resp(&mut self, ignore: bool) {
+        match self {
+            Self::H1(s) => s.set_ignore_info_resp(ignore),
+            Self::H2(_) => {} // always ignored
+        }
+    }
+
     /// Return a digest of the request including the method, path and Host header
     // TODO: make this use a `Formatter`
     pub fn request_summary(&self) -> String {
diff --git a/pingora-core/src/protocols/http/v1/common.rs b/pingora-core/src/protocols/http/v1/common.rs
index 18fa6c0..d6ade98 100644
--- a/pingora-core/src/protocols/http/v1/common.rs
+++ b/pingora-core/src/protocols/http/v1/common.rs
@@ -153,6 +153,14 @@ pub(super) fn is_upgrade_req(req: &RequestHeader) -> bool {
     req.version == http::Version::HTTP_11 && req.headers.get(header::UPGRADE).is_some()
 }
 
+pub(super) fn is_expect_continue_req(req: &RequestHeader) -> bool {
+    req.version == http::Version::HTTP_11
+        // https://www.rfc-editor.org/rfc/rfc9110#section-10.1.1
+        && req.headers.get(header::EXPECT).map_or(false, |v| {
+            v.as_bytes().eq_ignore_ascii_case(b"100-continue")
+        })
+}
+
 // Unlike the upgrade check on request, this function doesn't check the Upgrade or Connection header
 // because when seeing 101, we assume the server accepts to switch protocol.
 // In reality it is not common that some servers don't send all the required headers to establish
diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs
index 82a93a2..ea99d74 100644
--- a/pingora-core/src/protocols/http/v1/server.rs
+++ b/pingora-core/src/protocols/http/v1/server.rs
@@ -74,6 +74,8 @@ pub struct HttpSession {
     digest: Box<Digest>,
     /// Minimum send rate to the client
     min_send_rate: Option<usize>,
+    /// When this is enabled informational response headers will not be proxied downstream
+    ignore_info_resp: bool,
 }
 
 impl HttpSession {
@@ -109,6 +111,7 @@ impl HttpSession {
             upgraded: false,
             digest,
             min_send_rate: None,
+            ignore_info_resp: false,
         }
     }
 
@@ -388,6 +391,11 @@ impl HttpSession {
     /// Write the response header to the client.
     /// This function can be called more than once to send 1xx informational headers excluding 101.
     pub async fn write_response_header(&mut self, mut header: Box<ResponseHeader>) -> Result<()> {
+        if header.status.is_informational() && self.ignore_info_resp(header.status.into()) {
+            debug!("ignoring informational headers");
+            return Ok(());
+        }
+
         if let Some(resp) = self.response_written.as_ref() {
             if !resp.status.is_informational() || self.upgraded {
                 warn!("Respond header is already sent, cannot send again");
@@ -409,7 +417,7 @@ impl HttpSession {
             header.insert_header(header::CONNECTION, connection_value)?;
         }
 
-        if header.status.as_u16() == 101 {
+        if header.status == 101 {
             // make sure the connection is closed at the end when 101/upgrade is used
             self.set_keepalive(None);
         }
@@ -510,6 +518,18 @@ impl HttpSession {
         (None, None)
     }
 
+    fn ignore_info_resp(&self, status: u16) -> bool {
+        // ignore informational response if ignore flag is set and it's not an Upgrade and Expect: 100-continue isn't set
+        self.ignore_info_resp && status != 101 && !(status == 100 && self.is_expect_continue_req())
+    }
+
+    fn is_expect_continue_req(&self) -> bool {
+        match self.request_header.as_deref() {
+            Some(req) => is_expect_continue_req(req),
+            None => false,
+        }
+    }
+
     fn is_connection_keepalive(&self) -> Option<bool> {
         is_buf_keepalive(self.get_header(header::CONNECTION))
     }
@@ -824,6 +844,14 @@ impl HttpSession {
         }
     }
 
+    /// Sets whether we ignore writing informational responses downstream.
+    ///
+    /// This is a noop if the response is Upgrade or Continue and
+    /// Expect: 100-continue was set on the request.
+    pub fn set_ignore_info_resp(&mut self, ignore: bool) {
+        self.ignore_info_resp = ignore;
+    }
+
     /// Return the [Digest] of the connection.
     pub fn digest(&self) -> &Digest {
         &self.digest
@@ -1473,6 +1501,75 @@ mod tests_stream {
     }
 
     #[tokio::test]
+    async fn write_informational_ignored() {
+        let wire = b"HTTP/1.1 200 OK\r\nFoo: Bar\r\n\r\n";
+        let mock_io = Builder::new().write(wire).build();
+        let mut http_stream = HttpSession::new(Box::new(mock_io));
+        // ignore the 100 Continue
+        http_stream.ignore_info_resp = true;
+        let response_100 = ResponseHeader::build(StatusCode::CONTINUE, None).unwrap();
+        http_stream
+            .write_response_header_ref(&response_100)
+            .await
+            .unwrap();
+        let mut response_200 = ResponseHeader::build(StatusCode::OK, None).unwrap();
+        response_200.append_header("Foo", "Bar").unwrap();
+        http_stream.update_resp_headers = false;
+        http_stream
+            .write_response_header_ref(&response_200)
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn write_informational_100_not_ignored_if_expect_continue() {
+        let input = b"GET / HTTP/1.1\r\nExpect: 100-continue\r\n\r\n";
+        let output = b"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 200 OK\r\nFoo: Bar\r\n\r\n";
+
+        let mock_io = Builder::new().read(&input[..]).write(output).build();
+        let mut http_stream = HttpSession::new(Box::new(mock_io));
+        http_stream.read_request().await.unwrap();
+        http_stream.ignore_info_resp = true;
+        // 100 Continue is not ignored due to Expect: 100-continue on request
+        let response_100 = ResponseHeader::build(StatusCode::CONTINUE, None).unwrap();
+        http_stream
+            .write_response_header_ref(&response_100)
+            .await
+            .unwrap();
+        let mut response_200 = ResponseHeader::build(StatusCode::OK, None).unwrap();
+        response_200.append_header("Foo", "Bar").unwrap();
+        http_stream.update_resp_headers = false;
+        http_stream
+            .write_response_header_ref(&response_200)
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn write_informational_1xx_ignored_if_expect_continue() {
+        let input = b"GET / HTTP/1.1\r\nExpect: 100-continue\r\n\r\n";
+        let output = b"HTTP/1.1 200 OK\r\nFoo: Bar\r\n\r\n";
+
+        let mock_io = Builder::new().read(&input[..]).write(output).build();
+        let mut http_stream = HttpSession::new(Box::new(mock_io));
+        http_stream.read_request().await.unwrap();
+        http_stream.ignore_info_resp = true;
+        // 102 Processing is ignored
+        let response_102 = ResponseHeader::build(StatusCode::PROCESSING, None).unwrap();
+        http_stream
+            .write_response_header_ref(&response_102)
+            .await
+            .unwrap();
+        let mut response_200 = ResponseHeader::build(StatusCode::OK, None).unwrap();
+        response_200.append_header("Foo", "Bar").unwrap();
+        http_stream.update_resp_headers = false;
+        http_stream
+            .write_response_header_ref(&response_200)
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test]
     async fn write_101_switching_protocol() {
         let wire = b"HTTP/1.1 101 Switching Protocols\r\nFoo: Bar\r\n\r\n";
         let wire_body = b"nPAYLOAD";
diff --git a/pingora-core/src/protocols/http/v2/server.rs b/pingora-core/src/protocols/http/v2/server.rs
index f9072d6..5bd6bc0 100644
--- a/pingora-core/src/protocols/http/v2/server.rs
+++ b/pingora-core/src/protocols/http/v2/server.rs
@@ -194,22 +194,21 @@ impl HttpSession {
             return Ok(());
         }
 
-        // FIXME: we should ignore 1xx header because send_response() can only be called once
-        // https://github.com/hyperium/h2/issues/167
-
-        if let Some(resp) = self.response_written.as_ref() {
-            if !resp.status.is_informational() {
-                warn!("Respond header is already sent, cannot send again");
-                return Ok(());
-            }
+        if header.status.is_informational() {
+            // ignore informational response 1xx header because send_response() can only be called once
+            // https://github.com/hyperium/h2/issues/167
+            debug!("ignoring informational headers");
+            return Ok(());
         }
 
-        // no need to add these headers to 1xx responses
-        if !header.status.is_informational() {
-            /* update headers */
-            header.insert_header(header::DATE, get_cached_date())?;
+        if self.response_written.as_ref().is_some() {
+            warn!("Response header is already sent, cannot send again");
+            return Ok(());
         }
 
+        /* update headers */
+        header.insert_header(header::DATE, get_cached_date())?;
+
         // remove other h1 hop headers that cannot be present in H2
         // https://httpwg.org/specs/rfc7540.html#n-connection-specific-header-fields
         header.remove_header(&header::TRANSFER_ENCODING);
@@ -486,7 +485,8 @@ mod test {
         expected_trailers.insert("test", HeaderValue::from_static("trailers"));
         let trailers = expected_trailers.clone();
 
-        tokio::spawn(async move {
+        let mut handles = vec![];
+        handles.push(tokio::spawn(async move {
             let (h2, connection) = h2::client::handshake(client).await.unwrap();
             tokio::spawn(async move {
                 connection.await.unwrap();
@@ -510,7 +510,7 @@ mod test {
             assert_eq!(data, server_body);
             let resp_trailers = body.trailers().await.unwrap().unwrap();
             assert_eq!(resp_trailers, expected_trailers);
-        });
+        }));
 
         let mut connection = handshake(Box::new(server), None).await.unwrap();
         let digest = Arc::new(Digest::default());
@@ -520,7 +520,7 @@ mod test {
             .unwrap()
         {
             let trailers = trailers.clone();
-            tokio::spawn(async move {
+            handles.push(tokio::spawn(async move {
                 let req = http.req_header();
                 assert_eq!(req.method, Method::GET);
                 assert_eq!(req.uri, "https://www.example.com/");
@@ -545,7 +545,11 @@ mod test {
                 }
 
                 let response_header = Box::new(ResponseHeader::build(200, None).unwrap());
-                http.write_response_header(response_header, false).unwrap();
+                assert!(http
+                    .write_response_header(response_header.clone(), false)
+                    .is_ok());
+                // this write should be ignored otherwise we will error
+                assert!(http.write_response_header(response_header, false).is_ok());
 
                 // test idling after response header is sent
                 tokio::select! {
@@ -559,7 +563,11 @@ mod test {
 
                 http.write_trailers(trailers).unwrap();
                 http.finish().unwrap();
-            });
+            }));
+        }
+        for handle in handles {
+            // ensure no panics
+            assert!(handle.await.is_ok());
         }
     }
 }