Refactor `QPDFWriter` encryption logic to simplify pipeline writes and consolida…

…te encryption handling

Refactor `QPDFWriter` encryption logic to simplify pipeline writes and consolida…
…te encryption handling
m-holger
1 parent a78d18f7
Showing 3 changed files with 37 additions and 54 deletions
include/qpdf/QPDFWriter.hh
libqpdf/QPDFWriter.cc
libqpdf/qpdf/Pipeline_private.hh
@@ -466,6 +466,7 @@ class QPDFWriter
     QPDFWriter& write(std::integral auto val);
     QPDFWriter& write_name(std::string const& str);
     QPDFWriter& write_string(std::string const& str, bool force_binary = false);
+    QPDFWriter& write_encrypted(std::string_view str);
  
     template <typename... Args>
     QPDFWriter& write_qdf(Args&&... args);
@@ -569,7 +570,6 @@ class QPDFWriter
     // is popped.
  
     void adjustAESStreamLength(size_t& length);
-    PipelinePopper pushEncryptionFilter();
     void pushMD5Pipeline(PipelinePopper&);
     void computeDeterministicIDData();
  
@@ -1098,32 +1098,27 @@ QPDFWriter::adjustAESStreamLength(size_t&amp; length)
     }
 }
  
-QPDFWriter::PipelinePopper
-QPDFWriter::pushEncryptionFilter()
-{
-    auto pp = m->pipeline_stack.popper();
-    if (m->encryption && !m->cur_data_key.empty()) {
-        Pipeline* p = nullptr;
-        if (m->encrypt_use_aes) {
-            p = new Pl_AES_PDF(
-                "aes stream encryption",
-                m->pipeline,
+QPDFWriter&
+QPDFWriter::write_encrypted(std::string_view str)
+{
+    if (!(m->encryption && !m->cur_data_key.empty())) {
+        write(str);
+    } else if (m->encrypt_use_aes) {
+        write(
+            pl::pipe<Pl_AES_PDF>(
+                str,
                 true,
                 QUtil::unsigned_char_pointer(m->cur_data_key),
-                m->cur_data_key.length());
-        } else {
-            p = new Pl_RC4(
-                "rc4 stream encryption",
-                m->pipeline,
+                m->cur_data_key.length()));
+    } else {
+        write(
+            pl::pipe<Pl_RC4>(
+                str,
                 QUtil::unsigned_char_pointer(m->cur_data_key),
-                QIntC::to_int(m->cur_data_key.length()));
-        }
-        m->pipeline_stack.push(p);
+                QIntC::to_int(m->cur_data_key.length())));
     }
-    // Must call this unconditionally so we can call popPipelineStack to balance
-    // pushEncryptionFilter().
-    m->pipeline_stack.activate(pp);
-    return pp;
+
+    return *this;
 }
  
 void
@@ -1659,18 +1654,9 @@ QPDFWriter::unparseObject(
         adjustAESStreamLength(m->cur_stream_length);
         unparseObject(stream_dict, 0, flags, m->cur_stream_length, compress_stream);
         char last_char = stream_data.empty() ? '\0' : stream_data.back();
-        write("\nstream\n");
-        {
-            auto pp_enc = pushEncryptionFilter();
-            write(stream_data);
-        }
-
-        if ((m->added_newline =
-                 m->newline_before_endstream || (m->qdf_mode && last_char != '\n'))) {
-            write("\nendstream");
-        } else {
-            write("endstream");
-        }
+        write("\nstream\n").write_encrypted(stream_data);
+        m->added_newline = m->newline_before_endstream || (m->qdf_mode && last_char != '\n');
+        write(m->added_newline ? "\nendstream" : "endstream");
     } else if (tc == ::ot_string) {
         std::string val;
         if (m->encryption && !(flags & f_in_ostream) && !(flags & f_no_encryption) &&
@@ -1842,18 +1828,11 @@ QPDFWriter::writeObjectStream(QPDFObjectHandle object)
             unparseChild(extends, 1, f_in_ostream);
         }
     }
-    write_qdf("\n").write_no_qdf(" ").write(">>\nstream\n");
+    write_qdf("\n").write_no_qdf(" ").write(">>\nstream\n").write_encrypted(stream_buffer_pass2);
     if (m->encryption) {
         QTC::TC("qpdf", "QPDFWriter encrypt object stream");
     }
-    {
-        auto pp_enc = pushEncryptionFilter();
-        write(stream_buffer_pass2);
-    }
-    if (m->newline_before_endstream) {
-        write("\n");
-    }
-    write("endstream");
+    write(m->newline_before_endstream ? "\nendstream" : "endstream");
     m->cur_data_key.clear();
     closeObject(new_stream_id);
 }
@@ -2462,21 +2441,13 @@ QPDFWriter::writeHintStream(int hint_id)
     }
     adjustAESStreamLength(hlen);
     write(" /Length ").write(hlen);
-    write(" >>\nstream\n");
+    write(" >>\nstream\n").write_encrypted(hint_buffer);
  
     if (m->encryption) {
         QTC::TC("qpdf", "QPDFWriter encrypted hint stream");
     }
-    char last_char = hint_buffer.empty() ? '\0' : hint_buffer.back();
-    {
-        auto pp_enc = pushEncryptionFilter();
-        write(hint_buffer);
-    }
  
-    if (last_char != '\n') {
-        write("\n");
-    }
-    write("endstream");
+    write(hint_buffer.empty() || hint_buffer.back() != '\n' ? "\nendstream" : "endstream");
     closeObject(hint_id);
 }
  
@@ -173,6 +173,18 @@ namespace qpdf::pl
         unsigned long id_{0};
         bool pass_immediately_to_next{false};
     };
+
+    template <typename P, typename... Args>
+    std::string
+    pipe(std::string_view data, Args&&... args)
+    {
+        std::string result;
+        String s("", nullptr, result);
+        P pl("", &s, std::forward<Args>(args)...);
+        pl.write(reinterpret_cast<unsigned char const*>(data.data()), data.size());
+        pl.finish();
+        return result;
+    }
 } // namespace qpdf::pl
  
 #endif // PIPELINE_PRIVATE_HH