Refactor gzip decompression to share common code for chunked support.

PiperOrigin-RevId: 914468148

quiche/quic/masque/masque_ohttp_client.cc

diff --git a/quiche/quic/masque/masque_ohttp_client.cc b/quiche/quic/masque/masque_ohttp_client.cc
index 9c0c05a..e8d8056 100644
--- a/quiche/quic/masque/masque_ohttp_client.cc
+++ b/quiche/quic/masque/masque_ohttp_client.cc
@@ -104,51 +104,6 @@
   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)
@@ -246,6 +201,92 @@
 
 }  // namespace
 
+absl::StatusOr<GzipDecompressor::DecompressedData>
+GzipDecompressor::InflateLoop(int flush) {
+  int ret = Z_OK;
+  std::string decompressed = "";
+  std::vector<unsigned char> outbuffer(kGzipDecompressBufferSize);
+
+  // Decompress the input in chunks until we reach the end of the stream or
+  // need more input.
+  do {
+    stream_->next_out = outbuffer.data();
+    stream_->avail_out = outbuffer.size();
+
+    ret = inflate(stream_.get(), flush);
+
+    if (ret != Z_OK && ret != Z_STREAM_END) {
+      return absl::InternalError(absl::StrCat(
+          "Gzip decompression failed with error code: ", ret,
+          stream_->msg ? absl::StrCat(" (", stream_->msg, ")") : ""));
+    }
+
+    // Calculate how much data was placed in the buffer.
+    size_t decompressed_size = outbuffer.size() - stream_->avail_out;
+    decompressed.append(reinterpret_cast<char*>(outbuffer.data()),
+                        decompressed_size);
+  } while (stream_->avail_out == 0 && ret != Z_STREAM_END);
+
+  return DecompressedData{ret, std::move(decompressed)};
+}
+
+absl::StatusOr<std::unique_ptr<GzipDecompressor>> GzipDecompressor::Create() {
+  auto stream = std::make_unique<z_stream>();
+  // z_stream fields must be explicitly initialized, as per the zlib
+  // documentation.
+  memset(stream.get(), 0, sizeof(z_stream));
+
+  // 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`.
+  int inflate_result = inflateInit2(stream.get(), 16 + MAX_WBITS);
+  if (inflate_result != Z_OK) {
+    return absl::InternalError(
+        absl::StrCat("Failed to initialize zlib for gzip decompression, error "
+                     "code: ",
+                     inflate_result));
+  }
+
+  return std::unique_ptr<GzipDecompressor>(
+      new GzipDecompressor(std::move(stream)));
+}
+
+absl::StatusOr<std::string> GzipDecompressor::Decompress(
+    absl::string_view input, bool end_stream) {
+  if (stream_ == nullptr) {
+    return absl::FailedPreconditionError("Decompression has already ended");
+  }
+  stream_->next_in =
+      reinterpret_cast<unsigned char*>(const_cast<char*>(input.data()));
+  stream_->avail_in = input.size();
+
+  QUICHE_ASSIGN_OR_RETURN(DecompressedData decompressed_data,
+                          InflateLoop(end_stream ? Z_FINISH : Z_NO_FLUSH));
+
+  if (decompressed_data.status_code == Z_STREAM_END) {
+    finished_ = true;
+  }
+
+  if (end_stream && !finished_) {
+    return absl::InternalError(
+        "Gzip decompression failed to reach end of stream");
+  }
+  return decompressed_data.data;
+}
+
+void GzipDecompressor::EndDecompression() {
+  if (stream_ != nullptr) {
+    inflateEnd(stream_.get());
+    stream_.reset();
+  }
+}
+
+absl::StatusOr<std::string> GzipDecompress(absl::string_view input) {
+  QUICHE_ASSIGN_OR_RETURN(std::unique_ptr<GzipDecompressor> decompressor,
+                          GzipDecompressor::Create());
+  return decompressor->Decompress(input, /*end_stream=*/true);
+}
+
 std::string MasqueOhttpClient::Config::PerRequestConfig::method() const {
   if (method_.has_value()) {
     return *method_;

quiche/quic/masque/masque_ohttp_client.h

diff --git a/quiche/quic/masque/masque_ohttp_client.h b/quiche/quic/masque/masque_ohttp_client.h
index 132625e..6295b31 100644
--- a/quiche/quic/masque/masque_ohttp_client.h
+++ b/quiche/quic/masque/masque_ohttp_client.h
@@ -26,9 +26,55 @@
 #include "quiche/oblivious_http/buffers/oblivious_http_request.h"
 #include "quiche/oblivious_http/common/oblivious_http_chunk_handler.h"
 #include "quiche/oblivious_http/oblivious_http_client.h"
+#include <zlib.h>
 
 namespace quic {
 
+// A class that decompresses gzip encoded data using zlib. This also supports
+// incremental decompression, enabling the caller to feed a stream of compressed
+// chunks into the decompressor one by one.
+class QUICHE_NO_EXPORT GzipDecompressor {
+ public:
+  // Not copyable.
+  GzipDecompressor(const GzipDecompressor&) = delete;
+  GzipDecompressor& operator=(const GzipDecompressor&) = delete;
+
+  // Destructor is needed to free the internal state
+  // allocated by zlib during `inflateInit2`.
+  ~GzipDecompressor() { EndDecompression(); }
+
+  static absl::StatusOr<std::unique_ptr<GzipDecompressor>> Create();
+
+  // Can be called multiple times with continuous input chunks. `end_stream`
+  // should be set to true when the input is confirmed to have the final chunk.
+  absl::StatusOr<std::string> Decompress(absl::string_view input,
+                                         bool end_stream);
+
+  bool IsFinished() const { return finished_; }
+
+  // This will free the internal state allocated by zlib. The decompressor
+  // should not be used after this is called.
+  void EndDecompression();
+
+ private:
+  struct DecompressedData {
+    int status_code;
+    std::string data;
+  };
+
+  explicit GzipDecompressor(std::unique_ptr<z_stream> stream)
+      : stream_(std::move(stream)) {}
+
+  // Decompresses data from the zlib stream.
+  // `flush` determines the flushing behavior for zlib (e.g., Z_NO_FLUSH for
+  // chunks, Z_FINISH for single non-chunked). Returns the decompressed data and
+  // the zlib return code (e.g., Z_OK, Z_STREAM_END).
+  absl::StatusOr<DecompressedData> InflateLoop(int flush);
+
+  std::unique_ptr<z_stream> stream_;
+  bool finished_ = false;
+};
+
 // A client that sends OHTTP requests through a relay/gateway to target URLs.
 class QUICHE_EXPORT MasqueOhttpClient
     : public quic::MasqueConnectionPool::Visitor {