Merge pull request #1495 from m-holger/writer

Simplify QPDFWriter pipeline management

Merge pull request #1495 from m-holger/writer
Simplify QPDFWriter pipeline management
m-holger · GitHub
2 parents a5ec9dc0 9faffb74
Showing 10 changed files with 60 additions and 155 deletions
libqpdf/Pl_AES_PDF.cc
libqpdf/Pl_RC4.cc
libqpdf/QPDFWriter.cc
libqpdf/QPDF_encryption.cc
libqpdf/QPDF_linearization.cc
libqpdf/qpdf/Pipeline_private.hh
libqpdf/qpdf/Pl_AES_PDF.hh
libqpdf/qpdf/Pl_RC4.hh
libtests/aes.cc
libtests/rc4.cc
@@ -9,25 +9,18 @@
  
 bool Pl_AES_PDF::use_static_iv = false;
  
-Pl_AES_PDF::Pl_AES_PDF(
-    char const* identifier,
-    Pipeline* next,
-    bool encrypt,
-    unsigned char const* key,
-    size_t key_bytes) :
+Pl_AES_PDF::Pl_AES_PDF(char const* identifier, Pipeline* next, bool encrypt, std::string key) :
     Pipeline(identifier, next),
+    key(key),
     crypto(QPDFCryptoProvider::getImpl()),
-    encrypt(encrypt),
-    key_bytes(key_bytes)
+    encrypt(encrypt)
 {
     if (!next) {
         throw std::logic_error("Attempt to create Pl_AES_PDF with nullptr as next");
     }
-    if (!(key_bytes == 32 || key_bytes == 16)) {
+    if (!(key.size() == 32 || key.size() == 16)) {
         throw std::runtime_error("unsupported key length");
     }
-    this->key = std::make_unique<unsigned char[]>(key_bytes);
-    std::memcpy(this->key.get(), key, key_bytes);
     std::memset(this->inbuf, 0, this->buf_size);
     std::memset(this->outbuf, 0, this->buf_size);
     std::memset(this->cbc_block, 0, this->buf_size);
@@ -170,7 +163,12 @@ Pl_AES_PDF::flush(bool strip_padding)
                 return_after_init = true;
             }
         }
-        crypto->rijndael_init(encrypt, key.get(), key_bytes, cbc_mode, cbc_block);
+        crypto->rijndael_init(
+            encrypt,
+            reinterpret_cast<const unsigned char*>(key.data()),
+            key.size(),
+            cbc_mode,
+            cbc_block);
         if (return_after_init) {
             return;
         }
@@ -2,15 +2,10 @@
  
 #include <qpdf/QUtil.hh>
  
-Pl_RC4::Pl_RC4(
-    char const* identifier,
-    Pipeline* next,
-    unsigned char const* key_data,
-    int key_len,
-    size_t out_bufsize) :
+Pl_RC4::Pl_RC4(char const* identifier, Pipeline* next, std::string key, size_t out_bufsize) :
     Pipeline(identifier, next),
     out_bufsize(out_bufsize),
-    rc4(key_data, key_len)
+    rc4(reinterpret_cast<unsigned char const*>(key.data()), static_cast<int>(key.size()))
 {
     if (!next) {
         throw std::logic_error("Attempt to create Pl_RC4 with nullptr as next");
@@ -141,20 +141,20 @@ namespace
         }
  
         void
-        activate(Popper& pp, std::unique_ptr<pl::Link> link)
+        activate(Popper& pp, std::unique_ptr<Pipeline> next)
         {
             count_buffer.clear();
-            activate(pp, false, &count_buffer, std::move(link));
+            activate(pp, false, &count_buffer, std::move(next));
         }
  
         Popper
         activate(
             bool discard = false,
             std::string* str = nullptr,
-            std::unique_ptr<pl::Link> link = nullptr)
+            std::unique_ptr<Pipeline> next = nullptr)
         {
             Popper pp{*this};
-            activate(pp, discard, str, std::move(link));
+            activate(pp, discard, str, std::move(next));
             return pp;
         }
  
@@ -163,11 +163,11 @@ namespace
             Popper& pp,
             bool discard = false,
             std::string* str = nullptr,
-            std::unique_ptr<pl::Link> link = nullptr)
+            std::unique_ptr<Pipeline> next = nullptr)
         {
             std::unique_ptr<pl::Count> c;
-            if (link) {
-                c = std::make_unique<pl::Count>(++last_id, count_buffer, std::move(link));
+            if (next) {
+                c = std::make_unique<pl::Count>(++last_id, count_buffer, std::move(next));
             } else if (discard) {
                 c = std::make_unique<pl::Count>(++last_id, nullptr);
             } else if (!str) {
@@ -1109,18 +1109,9 @@ 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()));
+        write(pl::pipe<Pl_AES_PDF>(str, true, m->cur_data_key));
     } else {
-        write(
-            pl::pipe<Pl_RC4>(
-                str,
-                QUtil::unsigned_char_pointer(m->cur_data_key),
-                QIntC::to_int(m->cur_data_key.length())));
+        write(pl::pipe<Pl_RC4>(str, m->cur_data_key));
     }
  
     return *this;
@@ -1654,12 +1645,7 @@ QPDFWriter::unparseObject(
             val = object.getStringValue();
             if (m->encrypt_use_aes) {
                 Pl_Buffer bufpl("encrypted string");
-                Pl_AES_PDF pl(
-                    "aes encrypt string",
-                    &bufpl,
-                    true,
-                    QUtil::unsigned_char_pointer(m->cur_data_key),
-                    m->cur_data_key.length());
+                Pl_AES_PDF pl("aes encrypt string", &bufpl, true, m->cur_data_key);
                 pl.writeString(val);
                 pl.finish();
                 val = QPDF_String(bufpl.getString()).unparse(true);
@@ -1768,7 +1754,6 @@ QPDFWriter::writeObjectStream(QPDFObjectHandle object)
         }
     }
     {
-        auto pp_ostream = m->pipeline_stack.popper();
         // Adjust offsets to skip over comment before first object
         first = offsets.at(0);
         for (auto& iter: offsets) {
@@ -1783,18 +1768,15 @@ QPDFWriter::writeObjectStream(QPDFObjectHandle object)
         }
  
         // Set up a stream to write the stream data into a buffer.
-        if (compressed) {
-            m->pipeline_stack.activate(
-                pp_ostream,
-                pl::create<Pl_Flate>(
-                    pl::create<pl::String>(stream_buffer_pass2), Pl_Flate::a_deflate));
-        } else {
-            m->pipeline_stack.activate(pp_ostream, stream_buffer_pass2);
-        }
+        auto pp_ostream = m->pipeline_stack.activate(stream_buffer_pass2);
+
         writeObjectStreamOffsets(offsets, first_obj);
         write(stream_buffer_pass1);
         stream_buffer_pass1.clear();
         stream_buffer_pass1.shrink_to_fit();
+        if (compressed) {
+            stream_buffer_pass2 = pl::pipe<Pl_Flate>(stream_buffer_pass2, Pl_Flate::a_deflate);
+        }
     }
  
     // Write the object
@@ -2528,20 +2510,7 @@ QPDFWriter::writeXRefStream(
     std::string xref_data;
     const bool compressed = m->compress_streams && !m->qdf_mode;
     {
-        auto pp_xref = m->pipeline_stack.popper();
-        if (compressed) {
-            m->pipeline_stack.clear_buffer();
-            auto link = pl::create<pl::String>(xref_data);
-            if (!skip_compression) {
-                // Write the stream dictionary for compression but don't actually compress.  This
-                // helps us with computation of padding for pass 1 of linearization.
-                link = pl::create<Pl_Flate>(std::move(link), Pl_Flate::a_deflate);
-            }
-            m->pipeline_stack.activate(
-                pp_xref, pl::create<Pl_PNGFilter>(std::move(link), Pl_PNGFilter::a_encode, esize));
-        } else {
-            m->pipeline_stack.activate(pp_xref, xref_data);
-        }
+        auto pp_xref = m->pipeline_stack.activate(xref_data);
  
         for (int i = first; i <= last; ++i) {
             QPDFXRefEntry& e = m->new_obj[i].xref;
@@ -2577,6 +2546,15 @@ QPDFWriter::writeXRefStream(
         }
     }
  
+    if (compressed) {
+        xref_data = pl::pipe<Pl_PNGFilter>(xref_data, Pl_PNGFilter::a_encode, esize);
+        if (!skip_compression) {
+            // Write the stream dictionary for compression but don't actually compress.  This
+            // helps us with computation of padding for pass 1 of linearization.
+            xref_data = pl::pipe<Pl_Flate>(xref_data, Pl_Flate::a_deflate);
+        }
+    }
+
     openObject(xref_id);
     write("<<").write_qdf("\n ").write(" /Type /XRef").write_qdf("\n ");
     write(" /Length ").write(xref_data.size());
@@ -2743,9 +2721,7 @@ QPDFWriter::writeLinearized()
                 lin_pass1_file = QUtil::safe_fopen(m->lin_pass1_filename.c_str(), "wb");
                 m->pipeline_stack.activate(
                     pp_pass1,
-                    std::make_unique<pl::Link>(
-                        nullptr,
-                        std::make_unique<Pl_StdioFile>("linearization pass1", lin_pass1_file)));
+                    std::make_unique<Pl_StdioFile>("linearization pass1", lin_pass1_file));
             } else {
                 m->pipeline_stack.activate(pp_pass1, true);
             }
@@ -230,8 +230,7 @@ process_with_aes(
     size_t iv_length = 0)
 {
     Pl_Buffer buffer("buffer");
-    Pl_AES_PDF aes(
-        "aes", &buffer, encrypt, QUtil::unsigned_char_pointer(key), QIntC::to_uint(key.length()));
+    Pl_AES_PDF aes("aes", &buffer, encrypt, key);
     if (iv) {
         aes.setIV(iv, iv_length);
     } else {
@@ -902,12 +901,7 @@ QPDF::decryptString(std::string&amp; str, QPDFObjGen og)
         if (use_aes) {
             QTC::TC("qpdf", "QPDF_encryption aes decode string");
             Pl_Buffer bufpl("decrypted string");
-            Pl_AES_PDF pl(
-                "aes decrypt string",
-                &bufpl,
-                false,
-                QUtil::unsigned_char_pointer(key),
-                key.length());
+            Pl_AES_PDF pl("aes decrypt string", &bufpl, false, key);
             pl.writeString(str);
             pl.finish();
             str = bufpl.getString();
@@ -1028,19 +1022,11 @@ QPDF::decryptStream(
     std::string key = getKeyForObject(encp, og, use_aes);
     if (use_aes) {
         QTC::TC("qpdf", "QPDF_encryption aes decode stream");
-        decrypt_pipeline = std::make_unique<Pl_AES_PDF>(
-            "AES stream decryption",
-            pipeline,
-            false,
-            QUtil::unsigned_char_pointer(key),
-            key.length());
+        decrypt_pipeline =
+            std::make_unique<Pl_AES_PDF>("AES stream decryption", pipeline, false, key);
     } else {
         QTC::TC("qpdf", "QPDF_encryption rc4 decode stream");
-        decrypt_pipeline = std::make_unique<Pl_RC4>(
-            "RC4 stream decryption",
-            pipeline,
-            QUtil::unsigned_char_pointer(key),
-            toI(key.length()));
+        decrypt_pipeline = std::make_unique<Pl_RC4>("RC4 stream decryption", pipeline, key);
     }
     pipeline = decrypt_pipeline.get();
 }
@@ -1753,21 +1753,18 @@ QPDF::generateHintStream(
  
     // Write the hint stream itself into a compressed memory buffer. Write through a counter so we
     // can get offsets.
-    std::string b;
-    auto c = compressed
-        ? std::make_unique<pl::Count>(
-              0, b, pl::create<Pl_Flate>(pl::create<pl::String>(hint_buffer), Pl_Flate::a_deflate))
-        : std::make_unique<pl::Count>(0, hint_buffer);
-
-    BitWriter w(c.get());
+    pl::Count c(0, hint_buffer);
+    BitWriter w(&c);
  
     writeHPageOffset(w);
-    S = toI(c->getCount());
+    S = toI(c.getCount());
     writeHSharedObject(w);
     O = 0;
     if (m->outline_hints.nobjects > 0) {
-        O = toI(c->getCount());
+        O = toI(c.getCount());
         writeHGeneric(w, m->outline_hints);
     }
-    c->finish();
+    if (compressed) {
+        hint_buffer = pl::pipe<Pl_Flate>(hint_buffer, Pl_Flate::a_deflate);
+    }
 }
@@ -8,35 +8,6 @@
  
 namespace qpdf::pl
 {
-    struct Link
-    {
-        Link(std::unique_ptr<Link> next_link, std::unique_ptr<Pipeline> next_pl) :
-            next_link(std::move(next_link)),
-            next_pl(std::move(next_pl))
-        {
-        }
-
-        std::unique_ptr<Link> next_link{nullptr};
-        std::unique_ptr<Pipeline> next_pl{nullptr};
-    };
-
-    template <typename P, typename... Args>
-    std::unique_ptr<Link>
-    create(Args&&... args)
-    {
-        return std::make_unique<Link>(
-            nullptr, std::make_unique<P>("", nullptr, std::forward<Args>(args)...));
-    }
-
-    template <typename P, typename... Args>
-    std::unique_ptr<Link>
-    create(std::unique_ptr<Link> link, Args&&... args)
-    {
-        auto* next = link->next_pl.get();
-        return std::make_unique<Link>(
-            std::move(link), std::make_unique<P>("", next, std::forward<Args>(args)...));
-    }
-
     class String final: public Pipeline
     {
       public:
@@ -83,20 +54,10 @@ namespace qpdf::pl
         {
         }
  
-        // Count the number of characters written. If 'next' is not set, the content written will be
-        // discarded.
-        Count(unsigned long id, std::unique_ptr<Link> link) :
-            Pipeline("", link ? link->next_pl.get() : nullptr),
-            link(std::move(link)),
-            id_(id),
-            pass_immediately_to_next(link)
-        {
-        }
-
         // Write to 'str'. If 'next' is set, 'str' will be written to 'next' when 'finish' is
         // called.
-        Count(unsigned long id, std::string& str, std::unique_ptr<Link> link = nullptr) :
-            Pipeline("", link ? link->next_pl.get() : nullptr),
+        Count(unsigned long id, std::string& str, std::unique_ptr<Pipeline> link = nullptr) :
+            Pipeline("", link.get()),
             str(&str),
             link(std::move(link)),
             id_(id)
@@ -169,7 +130,7 @@ namespace qpdf::pl
       private:
         qpdf_offset_t count{0};
         std::string* str{nullptr};
-        std::unique_ptr<Link> link{nullptr};
+        std::unique_ptr<Pipeline> link{nullptr};
         unsigned long id_{0};
         bool pass_immediately_to_next{false};
     };
@@ -11,12 +11,7 @@ class Pl_AES_PDF final: public Pipeline
 {
   public:
     // key should be a pointer to key_bytes bytes of data
-    Pl_AES_PDF(
-        char const* identifier,
-        Pipeline* next,
-        bool encrypt,
-        unsigned char const* key,
-        size_t key_bytes);
+    Pl_AES_PDF(char const* identifier, Pipeline* next, bool encrypt, std::string key);
     ~Pl_AES_PDF() final = default;
  
     void write(unsigned char const* data, size_t len) final;
@@ -42,13 +37,12 @@ class Pl_AES_PDF final: public Pipeline
     static unsigned int const buf_size = QPDFCryptoImpl::rijndael_buf_size;
     static bool use_static_iv;
  
+    std::string key;
     std::shared_ptr<QPDFCryptoImpl> crypto;
     bool encrypt;
     bool cbc_mode{true};
     bool first{true};
     size_t offset{0}; // offset into memory buffer
-    std::unique_ptr<unsigned char[]> key;
-    size_t key_bytes{0};
     unsigned char inbuf[buf_size];
     unsigned char outbuf[buf_size];
     unsigned char cbc_block[buf_size];
@@ -12,11 +12,7 @@ class Pl_RC4 final: public Pipeline
  
     // key_len of -1 means treat key_data as a null-terminated string
     Pl_RC4(
-        char const* identifier,
-        Pipeline* next,
-        unsigned char const* key_data,
-        int key_len = -1,
-        size_t out_bufsize = def_bufsize);
+        char const* identifier, Pipeline* next, std::string key, size_t out_bufsize = def_bufsize);
     ~Pl_RC4() final = default;
  
     void write(unsigned char const* data, size_t len) final;
@@ -84,8 +84,9 @@ main(int argc, char* argv[])
         key[i / 2] = static_cast<unsigned char>(val);
     }
  
+    std::string keystr(reinterpret_cast<char const*>(key), keylen);
     auto* out = new Pl_StdioFile("stdout", outfile);
-    auto* aes = new Pl_AES_PDF("aes_128_cbc", out, encrypt, key, keylen);
+    auto* aes = new Pl_AES_PDF("aes_128_cbc", out, encrypt, keystr);
     delete[] key;
     key = nullptr;
     if (!cbc_mode) {
@@ -58,7 +58,8 @@ main(int argc, char* argv[])
     FILE* outfile = QUtil::safe_fopen(outfilename, "wb");
     auto* out = new Pl_StdioFile("stdout", outfile);
     // Use a small buffer size (64) for testing
-    auto* rc4 = new Pl_RC4("rc4", out, key, QIntC::to_int(keylen), 64U);
+    std::string keystr(reinterpret_cast<char const*>(key), keylen);
+    auto* rc4 = new Pl_RC4("rc4", out, keystr, 64U);
     delete[] key;
  
     // 64 < buffer size < 512, buffer_size is not a power of 2 for testing