Correctly init body writer for upgraded requests

Before this change, the code tried to get the header from the stored
header which is only set after this function is called.

2 files changed, 22 insertions, 2 deletions

diff --git a/.bleep b/.bleep
index a5b2b10..39a4637 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-541a53e91e59e55dc42ea0ec37880b11edd6bfe9
\ No newline at end of file
+75bf189b8e2ce2ab5acead15f8db45485085e577
\ No newline at end of file
diff --git a/pingora-core/src/protocols/http/v1/client.rs b/pingora-core/src/protocols/http/v1/client.rs
index 7881e59..c8e0300 100644
--- a/pingora-core/src/protocols/http/v1/client.rs
+++ b/pingora-core/src/protocols/http/v1/client.rs
@@ -532,7 +532,7 @@ impl HttpSession {
     }
 
     fn init_req_body_writer(&mut self, header: &RequestHeader) {
-        if self.is_upgrade_req() {
+        if is_upgrade_req(header) {
             self.body_writer.init_http10();
         } else {
             self.init_body_writer_comm(&header.headers)
@@ -893,6 +893,26 @@ mod tests_stream {
     }
 
     #[tokio::test]
+    async fn init_body_for_upgraded_req() {
+        use crate::protocols::http::v1::body::BodyMode;
+
+        let wire =
+            b"GET / HTTP/1.1\r\nConnection: Upgrade\r\nUpgrade: WS\r\nContent-Length: 0\r\n\r\n";
+        let mock_io = Builder::new().write(wire).build();
+        let mut http_stream = HttpSession::new(Box::new(mock_io));
+        let mut new_request = RequestHeader::build("GET", b"/", None).unwrap();
+        new_request.insert_header("Connection", "Upgrade").unwrap();
+        new_request.insert_header("Upgrade", "WS").unwrap();
+        // CL is ignored when Upgrade presents
+        new_request.insert_header("Content-Length", "0").unwrap();
+        let _ = http_stream
+            .write_request_header(Box::new(new_request))
+            .await
+            .unwrap();
+        assert_eq!(http_stream.body_writer.body_mode, BodyMode::HTTP1_0(0));
+    }
+
+    #[tokio::test]
     async fn read_switching_protocol() {
         init_log();
         let input1 = b"HTTP/1.1 101 Continue\r\n\r\n";