Avoid interim response handling for 101 Switching Protocols responses in BalsaFrame.

When updating SimpleHttpClient to set_use_interim_headers_callback(true) in
cl/527686928, I found that the client would not receive WebSocket payloads as
body after receiving the 101 response in some cases. This was likely because of
the state resetting in [1]. The resetting is consistent with 100 Continue
behavior and will likely also be helpful with 103 Early Hints (and other
interim responses), but is not helpful for Upgrade/101 Switching Protocols.

This CL special-cases 101 Switching Protocols, excluding these responses from
other interim response handling (involving OnInterimHeaders()), because 101
represents the "final" headers on the current protocol.

The added unit test fails without this change (missing, headers callback aside,
the body callbacks. Balsa instead consumes the entire message but only outputs
the header portion. Perhaps the body was erroneously considered as the start of
"final" headers).

[1] http://google3/third_party/quiche/balsa/balsa_frame.cc;l=857;rcl=524360383

PiperOrigin-RevId: 529163242

quiche/balsa/balsa_frame.cc

diff --git a/quiche/balsa/balsa_frame.cc b/quiche/balsa/balsa_frame.cc
index ac5e018..592a4db 100644
--- a/quiche/balsa/balsa_frame.cc
+++ b/quiche/balsa/balsa_frame.cc
@@ -43,7 +43,8 @@
 
 namespace {
 
-const size_t kContinueStatusCode = 100;
+constexpr size_t kContinueStatusCode = 100;
+constexpr size_t kSwitchingProtocolsStatusCode = 101;
 
 constexpr absl::string_view kChunked = "chunked";
 constexpr absl::string_view kContentLength = "content-length";
@@ -850,9 +851,11 @@
     }
 
     if (use_interim_headers_callback_ &&
-        IsInterimResponse(headers_->parsed_response_code())) {
+        IsInterimResponse(headers_->parsed_response_code()) &&
+        headers_->parsed_response_code() != kSwitchingProtocolsStatusCode) {
       // Deliver headers from this interim response but reset everything else to
-      // prepare for the next set of headers.
+      // prepare for the next set of headers. Skip 101 Switching Protocols
+      // because these are considered final headers for the current protocol.
       visitor_->OnInterimHeaders(std::move(*headers_));
       Reset();
       checkpoint = message_start = message_current;

quiche/balsa/balsa_frame_test.cc

diff --git a/quiche/balsa/balsa_frame_test.cc b/quiche/balsa/balsa_frame_test.cc
index 13a704a..83619d4 100644
--- a/quiche/balsa/balsa_frame_test.cc
+++ b/quiche/balsa/balsa_frame_test.cc
@@ -3978,6 +3978,40 @@
   EXPECT_EQ(balsa_frame_.ErrorCode(), BalsaFrameEnums::BALSA_NO_ERROR);
 }
 
+TEST_F(HTTPBalsaFrameTest, SwitchingProtocols) {
+  const std::string headers =
+      "HTTP/1.1 101 Switching Protocols\r\n"
+      "\r\n";
+  const std::string body = "Bytes for the new protocol";
+  const std::string message = absl::StrCat(headers, body);
+
+  // Even with the interim headers callback set, the 101 response is delivered
+  // as final response headers.
+  balsa_frame_.set_is_request(false);
+  balsa_frame_.set_use_interim_headers_callback(true);
+
+  InSequence s;
+  EXPECT_CALL(visitor_mock_, ProcessHeaders);
+  EXPECT_CALL(visitor_mock_, HeaderDone());
+
+  ASSERT_EQ(balsa_frame_.ProcessInput(message.data(), message.size()),
+            headers.size())
+      << balsa_frame_.ErrorCode();
+  ASSERT_FALSE(balsa_frame_.Error());
+  EXPECT_EQ(headers_.parsed_response_code(), 101);
+
+  balsa_frame_.AllowArbitraryBody();
+
+  EXPECT_CALL(visitor_mock_, OnRawBodyInput("Bytes for the new protocol"));
+  EXPECT_CALL(visitor_mock_, OnBodyChunkInput("Bytes for the new protocol"));
+  EXPECT_CALL(visitor_mock_, MessageDone()).Times(0);
+
+  ASSERT_EQ(balsa_frame_.ProcessInput(body.data(), body.size()), body.size());
+  EXPECT_FALSE(balsa_frame_.MessageFullyRead());
+  EXPECT_FALSE(balsa_frame_.Error());
+  EXPECT_EQ(balsa_frame_.ErrorCode(), BalsaFrameEnums::BALSA_NO_ERROR);
+}
+
 TEST_F(HTTPBalsaFrameTest, Http09) {
   constexpr absl::string_view request = "GET /\r\n";