Add GetBufferedDataSize to OHTTP chunked gateway and indeterminate-length BHTTP decoder. PiperOrigin-RevId: 945068201
diff --git a/quiche/binary_http/binary_http_message.h b/quiche/binary_http/binary_http_message.h index bbcdc49..cdd7949 100644 --- a/quiche/binary_http/binary_http_message.h +++ b/quiche/binary_http/binary_http_message.h
@@ -256,6 +256,9 @@ // truncated messages. absl::Status Decode(absl::string_view data, bool end_stream); + // Returns the number of bytes that are buffered in the decoder. + size_t GetBufferedDataSize() const { return buffer_.size(); } + private: // Carries out the decode logic from the checkpoint. Returns // OutOfRangeError if there is not enough data to decode the current
diff --git a/quiche/binary_http/binary_http_message_test.cc b/quiche/binary_http/binary_http_message_test.cc index 4f8e6d4..b07f087 100644 --- a/quiche/binary_http/binary_http_message_test.cc +++ b/quiche/binary_http/binary_http_message_test.cc
@@ -794,6 +794,165 @@ QUICHE_EXPECT_OK(ExpectTruncatedTrailerSection(message_data)); } +TEST(IndeterminateLengthDecoder, GetBufferedDataSize) { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + + // Construct the full message. + std::string request_bytes; + ASSERT_TRUE(absl::HexStringToBytes( + absl::StrCat( + k2ByteFramingIndicator, kIndeterminateLengthEncodedRequestControlData, + kIndeterminateLengthEncodedRequestHeaders, k8ByteContentTerminator, + kIndeterminateLengthEncodedRequestBodyChunks, k4ByteContentTerminator, + kIndeterminateLengthEncodedRequestTrailers, kContentTerminator, + kPadding), + &request_bytes)); + + // 1. Initially empty. + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + + // 2. Send partial framing + control data (4 bytes). + QUICHE_EXPECT_OK(decoder.Decode(absl::string_view(request_bytes).substr(0, 4), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 4); + + // 3. Send rest of control data (27 bytes). + // Total sent: 31 bytes (Framing + Control Data). + // Control data is finished, transition to kHeader. Buffer should be 0. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(4, 27), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + + // 4. Send partial headers. + // First field is 64 bytes. Send first field + 10 bytes of second field = 74 + // bytes. Total sent in this step: 74 bytes. Buffer should contain the 10 + // bytes of partial second field. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(31, 74), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 10); + + // 5. Send rest of headers (13 bytes) + partial headers terminator (4 bytes). + // Total sent in this step: 17 bytes. + // Buffer should contain the 4 bytes of partial terminator. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(31 + 74, 17), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 4); + + // 6. Send rest of headers terminator (4 bytes). + // Total sent in this step: 4 bytes. + // Headers section finished, transition to kBody. Buffer should be 0. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(31 + 74 + 17, 4), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + + // 7. Send partial body chunks. + // Body chunks total 21 bytes (3 chunks of 7 bytes each). + // Send Chunk 1 (7 bytes) + Chunk 2 (7 bytes) + partial Chunk 3 (3 bytes) = 17 + // bytes. Buffer should contain the 3 bytes of partial Chunk 3. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(126, 17), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 3); + + // 8. Send rest of Chunk 3 (4 bytes) + body terminator (4 bytes). + // Total sent in this step: 8 bytes. + // Body section finished, transition to kTrailer. Buffer should be 0. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(126 + 17, 8), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + + // 9. Send partial trailers. + // Trailers total 32 bytes (2 trailers of 16 bytes each). + // Send Trailer 1 (16 bytes) + partial Trailer 2 (5 bytes) = 21 bytes. + // Buffer should contain the 5 bytes of partial Trailer 2. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(151, 21), + /*end_stream=*/false)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 5); + + // 10. Send rest of Trailer 2 (11 bytes) + terminator (1 byte) + padding (3 + // bytes) with end_stream=true. Total sent in this step: 15 bytes. Message + // finished, buffer should be 0. + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(151 + 21, 15), + /*end_stream=*/true)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); +} + +TEST(IndeterminateLengthDecoder, GetBufferedDataSizeOnTruncation) { + // Construct the full message. + std::string request_bytes; + ASSERT_TRUE(absl::HexStringToBytes( + absl::StrCat( + k2ByteFramingIndicator, kIndeterminateLengthEncodedRequestControlData, + kIndeterminateLengthEncodedRequestHeaders, k8ByteContentTerminator, + kIndeterminateLengthEncodedRequestBodyChunks, k4ByteContentTerminator, + kIndeterminateLengthEncodedRequestTrailers, kContentTerminator, + kPadding), + &request_bytes)); + + // 1. Valid Truncation: End after headers. + { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + // Send up to headers terminator (126 bytes). + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(0, 126), + /*end_stream=*/true)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + } + + // 2. Valid Truncation: End after body chunks. + { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + // Send up to body terminator (151 bytes). + QUICHE_EXPECT_OK( + decoder.Decode(absl::string_view(request_bytes).substr(0, 151), + /*end_stream=*/true)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + } + + // 3. Invalid Truncation: End in middle of control data. + { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + // Send partial control data (4 bytes). + EXPECT_THAT(decoder.Decode(absl::string_view(request_bytes).substr(0, 4), + /*end_stream=*/true), + test::StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + } + + // 4. Invalid Truncation: End in middle of headers. + { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + // Send up to partial headers (31 + 74 = 105 bytes). + EXPECT_THAT(decoder.Decode(absl::string_view(request_bytes).substr(0, 105), + /*end_stream=*/true), + test::StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + } + + // 5. Invalid Truncation: End in middle of body chunk. + { + RequestMessageSectionTestHandler handler; + BinaryHttpRequest::IndeterminateLengthDecoder decoder(&handler); + // Send up to partial body chunk (126 + 17 = 143 bytes). + EXPECT_THAT(decoder.Decode(absl::string_view(request_bytes).substr(0, 143), + /*end_stream=*/true), + test::StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_EQ(decoder.GetBufferedDataSize(), 0); + } +} + struct RequestIndeterminateLengthEncoderTestData { BinaryHttpRequest::ControlData control_data{"POST", "https", "google.com", "/hello"};
diff --git a/quiche/oblivious_http/oblivious_http_gateway.cc b/quiche/oblivious_http/oblivious_http_gateway.cc index fdd8a3d..4796790 100644 --- a/quiche/oblivious_http/oblivious_http_gateway.cc +++ b/quiche/oblivious_http/oblivious_http_gateway.cc
@@ -229,6 +229,7 @@ absl::Status status = DecryptRequestCheckpoint(end_stream); if (end_stream) { request_current_section_ = RequestMessageSection::kEnd; + request_buffer_.clear(); if (absl::IsOutOfRange(status)) { // OutOfRange only used internally for buffering, so return // InvalidArgument if this is the end of the stream.
diff --git a/quiche/oblivious_http/oblivious_http_gateway.h b/quiche/oblivious_http/oblivious_http_gateway.h index d3a1089..1206ea0 100644 --- a/quiche/oblivious_http/oblivious_http_gateway.h +++ b/quiche/oblivious_http/oblivious_http_gateway.h
@@ -119,6 +119,9 @@ // portion of the final chunk. absl::Status DecryptRequest(absl::string_view data, bool end_stream); + // Returns the number of bytes that are buffered in the gateway. + size_t GetBufferedDataSize() const { return request_buffer_.size(); } + // Encrypts the data as a single chunk. If `is_final_chunk` is true, the // response will be encoded with the final chunk indicator. absl::StatusOr<std::string> EncryptResponse(
diff --git a/quiche/oblivious_http/oblivious_http_gateway_test.cc b/quiche/oblivious_http/oblivious_http_gateway_test.cc index c8d59b5..7e278a2 100644 --- a/quiche/oblivious_http/oblivious_http_gateway_test.cc +++ b/quiche/oblivious_http/oblivious_http_gateway_test.cc
@@ -172,6 +172,62 @@ EXPECT_EQ(chunk_handler.GetChunkCount(), 3); } +TEST(ChunkedObliviousHttpGateway, TestGetBufferedDataSize) { + TestChunkHandler chunk_handler; + auto instance = CreateChunkedObliviousHttpGateway(chunk_handler); + + EXPECT_EQ(instance->GetBufferedDataSize(), 0); + + std::string encapsulated_request_bytes; + ASSERT_TRUE(absl::HexStringToBytes(kEncapsulatedChunkedRequest, + &encapsulated_request_bytes)); + + // 1. Send header (39 bytes). + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(0, 39), + /*end_stream=*/false)); + EXPECT_EQ(instance->GetBufferedDataSize(), 0); + + // 2. Send partial Chunk 1 (15 bytes of 29 bytes chunk). + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(39, 15), + /*end_stream=*/false)); + EXPECT_EQ(instance->GetBufferedDataSize(), 15); + + // 3. Send remaining 14 bytes of Chunk 1. + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(39 + 15, 14), + /*end_stream=*/false)); + EXPECT_EQ(instance->GetBufferedDataSize(), 0); + + // 4. Send partial Chunk 2 (10 bytes of 30 bytes chunk). + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(68, 10), + /*end_stream=*/false)); + EXPECT_EQ(instance->GetBufferedDataSize(), 10); + + // 5. Send remaining 20 bytes of Chunk 2. + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(68 + 10, 20), + /*end_stream=*/false)); + EXPECT_EQ(instance->GetBufferedDataSize(), 0); + + // 6. Send final chunk indicator + partial final chunk data (1 byte indicator + // + 5 bytes data). + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(98, 6), + /*end_stream=*/false)); + // Only the 5 bytes of final chunk data should be buffered, indicator is + // consumed. + EXPECT_EQ(instance->GetBufferedDataSize(), 5); + + // 7. Send remaining final chunk data (11 bytes) with end_stream=true. + QUICHE_EXPECT_OK(instance->DecryptRequest( + absl::string_view(encapsulated_request_bytes).substr(98 + 6, 11), + /*end_stream=*/true)); + EXPECT_EQ(instance->GetBufferedDataSize(), 0); +} + TEST(ChunkedObliviousHttpGateway, DecryptingAfterDoneReturnsInvalidArgument) { TestChunkHandler chunk_handler; auto instance = CreateChunkedObliviousHttpGateway(chunk_handler);