Rework of CPDF_Object writing.

Move writing logic into implementation of related clases.

Change-Id: If70dc418b352b562ee681ea34fa6595d6f52eee3
Reviewed-on: https://pdfium-review.googlesource.com/36350
Commit-Queue: Art Snake <art-snake@yandex-team.ru>
Reviewed-by: Tom Sepez <tsepez@chromium.org>

5 files changed, 31 insertions, 155 deletions

diff --git a/core/fpdfapi/edit/cpdf_creator.cpp b/core/fpdfapi/edit/cpdf_creator.cpp
index f5d8c0fe49..5d4fac478d 100644
--- a/core/fpdfapi/edit/cpdf_creator.cpp
+++ b/core/fpdfapi/edit/cpdf_creator.cpp
@@ -151,143 +151,20 @@ CPDF_Creator::CPDF_Creator(CPDF_Document* pDoc,
 
 CPDF_Creator::~CPDF_Creator() {}
 
-bool CPDF_Creator::WriteStream(const CPDF_Object* pStream, uint32_t objnum) {
-  CPDF_CryptoHandler* pCrypto =
-      pStream != m_pMetadata ? GetCryptoHandler() : nullptr;
-
-  CPDF_FlateEncoder encoder(pStream->AsStream(), pStream != m_pMetadata);
-  CPDF_Encryptor encryptor(pCrypto, objnum, encoder.GetSpan());
-  if (static_cast<uint32_t>(encoder.GetDict()->GetIntegerFor("Length")) !=
-      encryptor.GetSpan().size()) {
-    encoder.CloneDict();
-    encoder.GetClonedDict()->SetNewFor<CPDF_Number>(
-        "Length", static_cast<int>(encryptor.GetSpan().size()));
-  }
-
-  if (!WriteDirectObj(objnum, encoder.GetDict(), true) ||
-      !m_Archive->WriteString("stream\r\n")) {
-    return false;
-  }
-
-  // Allow for empty streams.
-  if (encryptor.GetSpan().size() > 0 &&
-      !m_Archive->WriteBlock(encryptor.GetSpan().data(),
-                             encryptor.GetSpan().size())) {
-    return false;
-  }
-
-  return m_Archive->WriteString("\r\nendstream");
-}
-
 bool CPDF_Creator::WriteIndirectObj(uint32_t objnum, const CPDF_Object* pObj) {
   if (!m_Archive->WriteDWord(objnum) || !m_Archive->WriteString(" 0 obj\r\n"))
     return false;
 
-  if (pObj->IsStream()) {
-    if (!WriteStream(pObj, objnum))
-      return false;
-  } else if (!WriteDirectObj(objnum, pObj, true)) {
+  std::unique_ptr<CPDF_Encryptor> encryptor;
+  if (GetCryptoHandler() && pObj != m_pEncryptDict)
+    encryptor = pdfium::MakeUnique<CPDF_Encryptor>(GetCryptoHandler(), objnum);
+
+  if (!pObj->WriteTo(m_Archive.get(), encryptor.get()))
     return false;
-  }
 
   return m_Archive->WriteString("\r\nendobj\r\n");
 }
 
-bool CPDF_Creator::WriteDirectObj(uint32_t objnum,
-                                  const CPDF_Object* pObj,
-                                  bool bEncrypt) {
-  switch (pObj->GetType()) {
-    case CPDF_Object::BOOLEAN:
-    case CPDF_Object::NAME:
-    case CPDF_Object::NULLOBJ:
-    case CPDF_Object::NUMBER:
-    case CPDF_Object::REFERENCE:
-      if (!pObj->WriteTo(m_Archive.get()))
-        return false;
-      break;
-
-    case CPDF_Object::STRING: {
-      ByteString str = pObj->GetString();
-      bool bHex = pObj->AsString()->IsHex();
-      if (!GetCryptoHandler() || !bEncrypt) {
-        if (!pObj->WriteTo(m_Archive.get()))
-          return false;
-        break;
-      }
-      CPDF_Encryptor encryptor(GetCryptoHandler(), objnum, str.AsRawSpan());
-      ByteString content = PDF_EncodeString(
-          ByteString(encryptor.GetSpan().data(), encryptor.GetSpan().size()),
-          bHex);
-      if (!m_Archive->WriteString(content.AsStringView()))
-        return false;
-      break;
-    }
-    case CPDF_Object::STREAM: {
-      CPDF_FlateEncoder encoder(pObj->AsStream(), true);
-      CPDF_Encryptor encryptor(GetCryptoHandler(), objnum, encoder.GetSpan());
-      if (static_cast<uint32_t>(encoder.GetDict()->GetIntegerFor("Length")) !=
-          encryptor.GetSpan().size()) {
-        encoder.CloneDict();
-        encoder.GetClonedDict()->SetNewFor<CPDF_Number>(
-            "Length", static_cast<int>(encryptor.GetSpan().size()));
-      }
-      if (!WriteDirectObj(objnum, encoder.GetDict(), true) ||
-          !m_Archive->WriteString("stream\r\n") ||
-          !m_Archive->WriteBlock(encryptor.GetSpan().data(),
-                                 encryptor.GetSpan().size()) ||
-          !m_Archive->WriteString("\r\nendstream")) {
-        return false;
-      }
-
-      break;
-    }
-    case CPDF_Object::ARRAY: {
-      if (!m_Archive->WriteString("["))
-        return false;
-
-      const CPDF_Array* p = pObj->AsArray();
-      for (size_t i = 0; i < p->GetCount(); i++) {
-        if (!WriteDirectObj(objnum, p->GetObjectAt(i), true))
-          return false;
-      }
-      if (!m_Archive->WriteString("]"))
-        return false;
-      break;
-    }
-    case CPDF_Object::DICTIONARY: {
-      if (!GetCryptoHandler() || pObj == m_pEncryptDict) {
-        if (!pObj->WriteTo(m_Archive.get()))
-          return false;
-        break;
-      }
-
-      if (!m_Archive->WriteString("<<"))
-        return false;
-
-      const CPDF_Dictionary* p = pObj->AsDictionary();
-      bool bSignDict = p->IsSignatureDict();
-      for (const auto& it : *p) {
-        bool bSignValue = false;
-        const ByteString& key = it.first;
-        CPDF_Object* pValue = it.second.get();
-        if (!m_Archive->WriteString("/") ||
-            !m_Archive->WriteString(PDF_NameEncode(key).AsStringView())) {
-          return false;
-        }
-
-        if (bSignDict && key == "Contents")
-          bSignValue = true;
-        if (!WriteDirectObj(objnum, pValue, !bSignValue))
-          return false;
-      }
-      if (!m_Archive->WriteString(">>"))
-        return false;
-      break;
-    }
-  }
-  return true;
-}
-
 bool CPDF_Creator::WriteOldIndirectObject(uint32_t objnum) {
   if (m_pParser->IsObjectFreeOrNull(objnum))
     return true;
@@ -583,7 +460,7 @@ CPDF_Creator::Stage CPDF_Creator::WriteDoc_Stage4() {
           !m_Archive->WriteString(PDF_NameEncode(key).AsStringView())) {
         return Stage::kInvalid;
       }
-      if (!pValue->WriteTo(m_Archive.get()))
+      if (!pValue->WriteTo(m_Archive.get(), nullptr))
         return Stage::kInvalid;
     }
   } else {
@@ -632,7 +509,7 @@ CPDF_Creator::Stage CPDF_Creator::WriteDoc_Stage4() {
   }
   if (m_pIDArray) {
     if (!m_Archive->WriteString(("/ID")) ||
-        !m_pIDArray->WriteTo(m_Archive.get())) {
+        !m_pIDArray->WriteTo(m_Archive.get(), nullptr)) {
       return Stage::kInvalid;
     }
   }
diff --git a/core/fpdfapi/edit/cpdf_creator.h b/core/fpdfapi/edit/cpdf_creator.h
index a7261145e9..39e0950af6 100644
--- a/core/fpdfapi/edit/cpdf_creator.h
+++ b/core/fpdfapi/edit/cpdf_creator.h
@@ -69,11 +69,8 @@ class CPDF_Creator {
   bool WriteOldIndirectObject(uint32_t objnum);
   bool WriteOldObjs();
   bool WriteNewObjs();
-  bool WriteDirectObj(uint32_t objnum, const CPDF_Object* pObj, bool bEncrypt);
   bool WriteIndirectObj(uint32_t objnum, const CPDF_Object* pObj);
 
-  bool WriteStream(const CPDF_Object* pStream, uint32_t objnum);
-
   CPDF_CryptoHandler* GetCryptoHandler();
 
   UnownedPtr<CPDF_Document> const m_pDocument;
diff --git a/core/fpdfapi/edit/cpdf_encryptor.cpp b/core/fpdfapi/edit/cpdf_encryptor.cpp
index 8994b555ee..c74e53ed75 100644
--- a/core/fpdfapi/edit/cpdf_encryptor.cpp
+++ b/core/fpdfapi/edit/cpdf_encryptor.cpp
@@ -5,25 +5,27 @@
 // Original code copyright 2014 Foxit Software Inc. http://www.foxitsoftware.com
 
 #include "core/fpdfapi/edit/cpdf_encryptor.h"
+
 #include "core/fpdfapi/parser/cpdf_crypto_handler.h"
+#include "third_party/base/ptr_util.h"
 
-CPDF_Encryptor::CPDF_Encryptor(CPDF_CryptoHandler* pHandler,
-                               int objnum,
-                               pdfium::span<const uint8_t> src_data) {
-  if (src_data.empty())
-    return;
+CPDF_Encryptor::CPDF_Encryptor(CPDF_CryptoHandler* pHandler, int objnum)
+    : m_pHandler(pHandler), m_ObjNum(objnum) {
+  ASSERT(m_pHandler);
+}
 
-  if (!pHandler) {
-    m_Span = src_data;
-    return;
-  }
+std::vector<uint8_t> CPDF_Encryptor::Encrypt(
+    pdfium::span<const uint8_t> src_data) const {
+  if (src_data.empty())
+    return std::vector<uint8_t>();
 
-  uint32_t buf_size = pHandler->EncryptGetSize(src_data);
-  m_NewBuf.resize(buf_size);
-  pHandler->EncryptContent(objnum, 0, src_data, m_NewBuf.data(),
-                           buf_size);  // Updates |buf_size| with actual.
-  m_NewBuf.resize(buf_size);
-  m_Span = m_NewBuf;
+  std::vector<uint8_t> result;
+  uint32_t buf_size = m_pHandler->EncryptGetSize(src_data);
+  result.resize(buf_size);
+  m_pHandler->EncryptContent(m_ObjNum, 0, src_data, result.data(),
+                             buf_size);  // Updates |buf_size| with actual.
+  result.resize(buf_size);
+  return result;
 }
 
 CPDF_Encryptor::~CPDF_Encryptor() {}
diff --git a/core/fpdfapi/edit/cpdf_encryptor.h b/core/fpdfapi/edit/cpdf_encryptor.h
index 67fd5c4f9b..5486c4f248 100644
--- a/core/fpdfapi/edit/cpdf_encryptor.h
+++ b/core/fpdfapi/edit/cpdf_encryptor.h
@@ -9,25 +9,25 @@
 
 #include <stdint.h>
 
+#include <memory>
 #include <vector>
 
 #include "core/fxcrt/fx_memory.h"
+#include "core/fxcrt/unowned_ptr.h"
 #include "third_party/base/span.h"
 
 class CPDF_CryptoHandler;
 
 class CPDF_Encryptor {
  public:
-  CPDF_Encryptor(CPDF_CryptoHandler* pHandler,
-                 int objnum,
-                 pdfium::span<const uint8_t> src_data);
+  CPDF_Encryptor(CPDF_CryptoHandler* pHandler, int objnum);
   ~CPDF_Encryptor();
 
-  pdfium::span<const uint8_t> GetSpan() const { return m_Span; }
+  std::vector<uint8_t> Encrypt(pdfium::span<const uint8_t> src_data) const;
 
  private:
-  std::vector<uint8_t> m_NewBuf;
-  pdfium::span<const uint8_t> m_Span;
+  UnownedPtr<CPDF_CryptoHandler> const m_pHandler;
+  const int m_ObjNum;
 };
 
 #endif  // CORE_FPDFAPI_EDIT_CPDF_ENCRYPTOR_H_
diff --git a/core/fpdfapi/edit/cpdf_pagecontentgenerator.cpp b/core/fpdfapi/edit/cpdf_pagecontentgenerator.cpp
index f6a941d200..9693bc48ff 100644
--- a/core/fpdfapi/edit/cpdf_pagecontentgenerator.cpp
+++ b/core/fpdfapi/edit/cpdf_pagecontentgenerator.cpp
@@ -266,7 +266,7 @@ const CPDF_ContentMark* CPDF_PageContentGenerator::ProcessContentMarks(
     // If there are parameters, write properties, direct or indirect.
     if (item->GetParamType() == CPDF_ContentMarkItem::DirectDict) {
       CPDF_StringArchiveStream archive_stream(buf);
-      item->GetParam()->WriteTo(&archive_stream);
+      item->GetParam()->WriteTo(&archive_stream, nullptr);
       *buf << " ";
     } else {
       ASSERT(item->GetParamType() == CPDF_ContentMarkItem::PropertiesDict);