From 2bb9e06d1eb779edfe17f8ce37871ee3b1959211 Mon Sep 17 00:00:00 2001
From: m-holger <m.holger@qpdf.org>
Date: Sun, 28 Jul 2024 19:53:46 +0100
Subject: [PATCH] In qpdf_fuzzer add a memory limit for Pl_Flate

---
 fuzz/qpdf_fuzzer.cc      |  2 ++
 include/qpdf/Pl_Flate.hh |  6 ++++++
 libqpdf/Pl_Flate.cc      | 17 +++++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/fuzz/qpdf_fuzzer.cc b/fuzz/qpdf_fuzzer.cc
index 40504af0..f41933dc 100644
--- a/fuzz/qpdf_fuzzer.cc
+++ b/fuzz/qpdf_fuzzer.cc
@@ -2,6 +2,7 @@
 #include <qpdf/BufferInputSource.hh>
 #include <qpdf/Pl_DCT.hh>
 #include <qpdf/Pl_Discard.hh>
+#include <qpdf/Pl_Flate.hh>
 #include <qpdf/Pl_PNGFilter.hh>
 #include <qpdf/Pl_TIFFPredictor.hh>
 #include <qpdf/QPDF.hh>
@@ -183,6 +184,7 @@ FuzzHelper::doChecks()
 
     Pl_PNGFilter::setMemoryLimit(1'000'000);
     Pl_TIFFPredictor::setMemoryLimit(1'000'000);
+    Pl_Flate::setMemoryLimit(10'000'000);
 
     // Do not decompress corrupt data. This may cause extended runtime within jpeglib without
     // exercising additional code paths in qpdf, and potentially causing counterproductive timeouts.
diff --git a/include/qpdf/Pl_Flate.hh b/include/qpdf/Pl_Flate.hh
index 41ae7e4a..bb5de2ab 100644
--- a/include/qpdf/Pl_Flate.hh
+++ b/include/qpdf/Pl_Flate.hh
@@ -42,6 +42,11 @@ class QPDF_DLL_CLASS Pl_Flate: public Pipeline
     QPDF_DLL
     ~Pl_Flate() override;
 
+    // Limit the memory used.
+    // NB This is a static option affecting all Pl_PNGFilter instances.
+    QPDF_DLL
+    static void setMemoryLimit(unsigned long long limit);
+
     QPDF_DLL
     void write(unsigned char const* data, size_t len) override;
     QPDF_DLL
@@ -87,6 +92,7 @@ class QPDF_DLL_CLASS Pl_Flate: public Pipeline
         action_e action;
         bool initialized;
         void* zdata;
+        unsigned long long written{0};
         std::function<void(char const*, int)> callback;
     };
 
diff --git a/libqpdf/Pl_Flate.cc b/libqpdf/Pl_Flate.cc
index d332e635..dd50f902 100644
--- a/libqpdf/Pl_Flate.cc
+++ b/libqpdf/Pl_Flate.cc
@@ -7,6 +7,11 @@
 #include <qpdf/QIntC.hh>
 #include <qpdf/QUtil.hh>
 
+namespace
+{
+    unsigned long long memory_limit{0};
+} // namespace
+
 int Pl_Flate::compression_level = Z_DEFAULT_COMPRESSION;
 
 Pl_Flate::Members::Members(size_t out_bufsize, action_e action) :
@@ -63,6 +68,12 @@ Pl_Flate::~Pl_Flate() // NOLINT (modernize-use-equals-default)
     // Must be explicit and not inline -- see QPDF_DLL_CLASS in README-maintainer
 }
 
+void
+Pl_Flate::setMemoryLimit(unsigned long long limit)
+{
+    memory_limit = limit;
+}
+
 void
 Pl_Flate::setWarnCallback(std::function<void(char const*, int)> callback)
 {
@@ -170,6 +181,12 @@ Pl_Flate::handleData(unsigned char const* data, size_t len, int flush)
                 }
                 uLong ready = QIntC::to_ulong(m->out_bufsize - zstream.avail_out);
                 if (ready > 0) {
+                    if (memory_limit) {
+                        m->written += ready;
+                        if (m->written > memory_limit) {
+                            throw std::runtime_error("PL_Flate memory limit exceeded");
+                        }
+                    }
                     this->getNext()->write(m->outbuf.get(), ready);
                     zstream.next_out = m->outbuf.get();
                     zstream.avail_out = QIntC::to_uint(m->out_bufsize);