Ignore write_response_header on upgraded websockets

cloudflare · May 24, 2024 · cd51344 · cd51344
1 parent a401754
commit cd51344
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/.bleep b/.bleep
@@ -1 +1 @@
-7761fa09483329f00ba1945076e27abb53bc1f26
+d44d3b13cb33c5163561923a94d62baed216767d
diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs
@@ -386,7 +386,7 @@ impl HttpSession {
     /// 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 let Some(resp) = self.response_written.as_ref() {
-            if !resp.status.is_informational() {
+            if !resp.status.is_informational() || self.upgraded {
                 warn!("Respond header is already sent, cannot send again");
                 return Ok(());
             }
@@ -1431,6 +1431,14 @@ mod tests_stream {
             .unwrap();
         let n = http_stream.write_body(wire_body).await.unwrap().unwrap();
         assert_eq!(wire_body.len(), n);
+        // simulate upgrade
+        http_stream.upgraded = true;
+        // this write should be ignored
+        let response_502 = ResponseHeader::build(StatusCode::BAD_GATEWAY, None).unwrap();
+        http_stream
+            .write_response_header_ref(&response_502)
+            .await
+            .unwrap();
     }
 
     #[tokio::test]
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		7761fa09483329f00ba1945076e27abb53bc1f26
		d44d3b13cb33c5163561923a94d62baed216767d