Add gzip response decompression support to Masque OHTTP client. This change adds the `--gzip` flag. When enabled, the client adds `Accept-Encoding: gzip` to inner requests and decompresses inner responses that have `Content-Encoding: gzip` using zlib. PiperOrigin-RevId: 907113338
diff --git a/quiche/BUILD.bazel b/quiche/BUILD.bazel index cfa8f97..06a507f 100644 --- a/quiche/BUILD.bazel +++ b/quiche/BUILD.bazel
@@ -437,6 +437,7 @@ "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@com_google_absl//absl/types:variant", + "@zlib", ], )
diff --git a/quiche/quic/masque/masque_ohttp_client.cc b/quiche/quic/masque/masque_ohttp_client.cc index 5acd4d8..9c0c05a 100644 --- a/quiche/quic/masque/masque_ohttp_client.cc +++ b/quiche/quic/masque/masque_ohttp_client.cc
@@ -4,6 +4,7 @@ #include "quiche/quic/masque/masque_ohttp_client.h" +#include <cstring> #include <functional> #include <iostream> #include <memory> @@ -42,6 +43,7 @@ #include "quiche/oblivious_http/buffers/oblivious_http_response.h" #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h" #include "quiche/oblivious_http/oblivious_http_client.h" +#include <zlib.h> namespace quic { @@ -61,6 +63,8 @@ namespace { static constexpr uint64_t kFixedSizeResponseFramingIndicator = 0x01; +// Buffer size for gzip decompression. 32KB chosen arbitrarily. +static constexpr size_t kGzipDecompressBufferSize = 32768; absl::Status ParseHeadersIntoMap( const std::vector<std::string>& headers, @@ -100,6 +104,51 @@ return chunks; } +absl::StatusOr<std::string> GzipDecompress(absl::string_view input) { + z_stream zs; + memset(&zs, 0, sizeof(zs)); + + // Initialize zlib for gzip decompression. + // 16 + MAX_WBITS tells zlib to expect a gzip header and trailer, which is the + // expectation for valid HTTP responses with `Content-Encoding: gzip`. + if (inflateInit2(&zs, 16 + MAX_WBITS) != Z_OK) { + return absl::InternalError( + "Failed to initialize zlib for gzip decompression"); + } + + // Automatically clean up zlib resources when exiting the function. + absl::Cleanup cleanup = [&zs] { inflateEnd(&zs); }; + + zs.next_in = reinterpret_cast<uint8_t*>(const_cast<char*>(input.data())); + zs.avail_in = input.size(); + + int ret; + std::vector<uint8_t> outbuffer(kGzipDecompressBufferSize); + std::string decompressed; + + // Decompress the input in chunks until we reach the end of the stream. + do { + zs.next_out = outbuffer.data(); + zs.avail_out = outbuffer.size(); + + ret = inflate(&zs, Z_NO_FLUSH); + + // Calculate how much data was placed in the buffer and append it. + size_t decompressed_size = outbuffer.size() - zs.avail_out; + decompressed.append(reinterpret_cast<char*>(outbuffer.data()), + decompressed_size); + } while (ret == Z_OK); + + // Z_STREAM_END indicates that the full compressed stream was processed + // successfully. + if (ret != Z_STREAM_END) { + return absl::InternalError( + absl::StrCat("Gzip decompression failed with error code: ", ret)); + } + + return decompressed; +} + class PingPongResponseVisitor : public MasqueOhttpClient::ResponseVisitor { public: explicit PingPongResponseVisitor(std::vector<std::string> chunks) @@ -544,6 +593,9 @@ per_request_config.headers()) { headers.push_back({header.first, header.second}); } + if (config_.handle_gzip_response()) { + headers.push_back({"accept-encoding", "gzip"}); + } QUICHE_ASSIGN_OR_RETURN( std::string encoded_headers, pending_request.encoder->EncodeHeaders(absl::MakeSpan(headers))); @@ -574,6 +626,9 @@ per_request_config.headers()) { binary_request.AddHeaderField({header.first, header.second}); } + if (config_.handle_gzip_response()) { + binary_request.AddHeaderField({"accept-encoding", "gzip"}); + } binary_request.set_body(post_data); QUICHE_ASSIGN_OR_RETURN(encoded_data, binary_request.Serialize()); } @@ -795,6 +850,20 @@ << request_id << ". Body length is " << encapsulated_response->body.size() << ". Headers:" << encapsulated_response->headers.DebugString(); + if (config_.handle_gzip_response()) { + auto content_encoding_it = + encapsulated_response->headers.find("content-encoding"); + if (content_encoding_it != encapsulated_response->headers.end() && + absl::EqualsIgnoreCase(content_encoding_it->second, "gzip")) { + size_t compressed_size = encapsulated_response->body.size(); + QUICHE_ASSIGN_OR_RETURN(std::string decompressed_body, + GzipDecompress(encapsulated_response->body)); + QUICHE_LOG(INFO) << "Successfully decompressed gzip response from size " + << compressed_size << " to size " + << decompressed_body.size(); + encapsulated_response->body = std::move(decompressed_body); + } + } std::cout << encapsulated_response->body; int16_t encapsulated_status_code = MasqueConnectionPool::GetStatusCode(*encapsulated_response);
diff --git a/quiche/quic/masque/masque_ohttp_client.h b/quiche/quic/masque/masque_ohttp_client.h index 69669a8..132625e 100644 --- a/quiche/quic/masque/masque_ohttp_client.h +++ b/quiche/quic/masque/masque_ohttp_client.h
@@ -142,6 +142,9 @@ void SetDnsConfig(const MasqueConnectionPool::DnsConfig& dns_config) { dns_config_ = dns_config; } + void SetHandleGzipResponse(bool handle_gzip_response) { + handle_gzip_response_ = handle_gzip_response; + } absl::Status AddKeyFetchHeaders( const std::vector<std::string>& key_fetch_headers); void AddPerRequestConfig(const PerRequestConfig& per_request_config) { @@ -165,6 +168,7 @@ const { return key_fetch_headers_; } + bool handle_gzip_response() const { return handle_gzip_response_; } private: std::string key_fetch_url_; @@ -175,6 +179,7 @@ MasqueConnectionPool::DnsConfig dns_config_; std::vector<std::pair<std::string, std::string>> key_fetch_headers_; std::vector<PerRequestConfig> per_request_configs_; + bool handle_gzip_response_ = false; }; class ResponseVisitor {
diff --git a/quiche/quic/masque/masque_ohttp_client_bin.cc b/quiche/quic/masque/masque_ohttp_client_bin.cc index b82a1b8..4cf330e 100644 --- a/quiche/quic/masque/masque_ohttp_client_bin.cc +++ b/quiche/quic/masque/masque_ohttp_client_bin.cc
@@ -25,6 +25,11 @@ bool, disable_certificate_verification, false, "If true, don't verify the server certificate."); +DEFINE_QUICHE_COMMAND_LINE_FLAG(bool, handle_gzip_response, false, + "If true, adds `accept-encoding: gzip` to the " + "inner request and decompresses " + "gzip-encoded inner responses."); + DEFINE_QUICHE_COMMAND_LINE_FLAG( bool, use_mtls_for_key_fetch, false, "If true, use mTLS when fetching the OHTTP/HPKE keys."); @@ -126,6 +131,8 @@ const bool disable_certificate_verification = quiche::GetQuicheCommandLineFlag(FLAGS_disable_certificate_verification); + const bool handle_gzip_response = + quiche::GetQuicheCommandLineFlag(FLAGS_handle_gzip_response); const bool use_mtls_for_key_fetch = quiche::GetQuicheCommandLineFlag(FLAGS_use_mtls_for_key_fetch); const std::string client_cert_file = @@ -206,6 +213,7 @@ } MasqueOhttpClient::Config config(/*key_fetch_url=*/urls[0], /*relay_url=*/urls[1]); + config.SetHandleGzipResponse(handle_gzip_response); if (use_mtls_for_key_fetch) { QUICHE_RETURN_IF_ERROR(config.ConfigureKeyFetchClientCert( client_cert_file, client_cert_key_file));