Add Embeddertest for password after savingchromium/3157

This CL allows TestSaved to open password-protected and  linearized
documents properly. It also adds a test with one such document.

BUG= pdfium:787

Change-Id: Ie0da7f290711505fb208794afdc737c36e84dd3c
Reviewed-on: https://pdfium-review.googlesource.com/7034
Commit-Queue: Nicolás Peña <npm@chromium.org>
Reviewed-by: dsinclair <dsinclair@chromium.org>

3 files changed, 115 insertions, 40 deletions

diff --git a/core/fpdfapi/parser/cpdf_security_handler_embeddertest.cpp b/core/fpdfapi/parser/cpdf_security_handler_embeddertest.cpp
index c6c6217e2c..5df3348834 100644
--- a/core/fpdfapi/parser/cpdf_security_handler_embeddertest.cpp
+++ b/core/fpdfapi/parser/cpdf_security_handler_embeddertest.cpp
@@ -2,6 +2,12 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <string>
+
+#include "core/fxcrt/fx_system.h"
+#include "public/fpdf_edit.h"
+#include "public/fpdf_save.h"
+#include "public/fpdfview.h"
 #include "testing/embedder_test.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -35,6 +41,48 @@ TEST_F(CPDFSecurityHandlerEmbeddertest, OwnerPassword) {
   EXPECT_EQ(0xFFFFFFFC, FPDF_GetDocPermissions(document()));
 }
 
+TEST_F(CPDFSecurityHandlerEmbeddertest, PasswordAfterGenerateSave) {
+#if _FXM_PLATFORM_ == _FXM_PLATFORM_LINUX_
+  const char md5[] = "e4a3701ca5b2a759e06455aa8d97d46e";
+#elif _FXM_PLATFORM_ == _FXM_PLATFORM_APPLE_
+  const char md5[] = "6951b6c9891dfe0332a5b1983e484400";
+#else
+  const char md5[] = "50985f3440d3f66c3b599ab138214015";
+#endif  // _FXM_PLATFORM_ == _FXM_PLATFORM_LINUX_
+  {
+    ASSERT_TRUE(OpenDocument("encrypted.pdf", "5678", true));
+    FPDF_PAGE page = LoadPage(0);
+    ASSERT_TRUE(page);
+    FPDF_PAGEOBJECT red_rect = FPDFPageObj_CreateNewRect(10, 10, 20, 20);
+    ASSERT_TRUE(red_rect);
+    EXPECT_TRUE(FPDFPath_SetFillColor(red_rect, 255, 0, 0, 255));
+    EXPECT_TRUE(FPDFPath_SetDrawMode(red_rect, FPDF_FILLMODE_ALTERNATE, 0));
+    FPDFPage_InsertObject(page, red_rect);
+    FPDF_BITMAP page_bitmap = RenderPage(page);
+    CompareBitmap(page_bitmap, 612, 792, md5);
+    FPDFBitmap_Destroy(page_bitmap);
+    EXPECT_TRUE(FPDFPage_GenerateContent(page));
+    EXPECT_TRUE(FPDF_SaveAsCopy(document(), this, 0));
+    UnloadPage(page);
+  }
+  std::string new_file = GetString();
+  FPDF_FILEACCESS file_access;
+  memset(&file_access, 0, sizeof(file_access));
+  file_access.m_FileLen = new_file.size();
+  file_access.m_GetBlock = GetBlockFromString;
+  file_access.m_Param = &new_file;
+  EXPECT_FALSE(FPDF_LoadCustomDocument(&file_access, nullptr));
+  struct {
+    const char* password;
+    const unsigned long permissions;
+  } tests[] = {{"1234", 0xFFFFF2C0}, {"5678", 0xFFFFFFFC}};
+  for (const auto& test : tests) {
+    TestSaved(612, 792, md5, test.password);
+    EXPECT_EQ(test.permissions, FPDF_GetDocPermissions(m_SavedDocument));
+    CloseSaved();
+  }
+}
+
 TEST_F(CPDFSecurityHandlerEmbeddertest, NoPasswordVersion5) {
   ASSERT_FALSE(OpenDocument("bug_644.pdf"));
 }
diff --git a/testing/embedder_test.cpp b/testing/embedder_test.cpp
index 80d4a66381..0846d8ccb7 100644
--- a/testing/embedder_test.cpp
+++ b/testing/embedder_test.cpp
@@ -135,59 +135,67 @@ bool EmbedderTest::OpenDocument(const std::string& filename,
   file_access_.m_FileLen = static_cast<unsigned long>(file_length_);
   file_access_.m_GetBlock = TestLoader::GetBlock;
   file_access_.m_Param = loader_;
+  return OpenDocumentHelper(password, must_linearize, &file_avail_, &hints_,
+                            &file_access_, &document_, &avail_, &form_handle_);
+}
 
-  file_avail_.version = 1;
-  file_avail_.IsDataAvail = Is_Data_Avail;
-
-  hints_.version = 1;
-  hints_.AddSegment = Add_Segment;
-
-  avail_ = FPDFAvail_Create(&file_avail_, &file_access_);
-
-  if (FPDFAvail_IsLinearized(avail_) == PDF_LINEARIZED) {
-    document_ = FPDFAvail_GetDocument(avail_, password);
-    if (!document_) {
+bool EmbedderTest::OpenDocumentHelper(const char* password,
+                                      bool must_linearize,
+                                      FX_FILEAVAIL* file_avail,
+                                      FX_DOWNLOADHINTS* hints,
+                                      FPDF_FILEACCESS* file_access,
+                                      FPDF_DOCUMENT* document,
+                                      FPDF_AVAIL* avail,
+                                      FPDF_FORMHANDLE* form_handle) {
+  file_avail->version = 1;
+  file_avail->IsDataAvail = Is_Data_Avail;
+
+  hints->version = 1;
+  hints->AddSegment = Add_Segment;
+
+  *avail = FPDFAvail_Create(file_avail, file_access);
+
+  if (FPDFAvail_IsLinearized(*avail) == PDF_LINEARIZED) {
+    *document = FPDFAvail_GetDocument(*avail, password);
+    if (!*document)
       return false;
-    }
+
     int32_t nRet = PDF_DATA_NOTAVAIL;
-    while (nRet == PDF_DATA_NOTAVAIL) {
-      nRet = FPDFAvail_IsDocAvail(avail_, &hints_);
-    }
-    if (nRet == PDF_DATA_ERROR) {
+    while (nRet == PDF_DATA_NOTAVAIL)
+      nRet = FPDFAvail_IsDocAvail(*avail, hints);
+    if (nRet == PDF_DATA_ERROR)
       return false;
-    }
-    nRet = FPDFAvail_IsFormAvail(avail_, &hints_);
-    if (nRet == PDF_FORM_ERROR || nRet == PDF_FORM_NOTAVAIL) {
+
+    nRet = FPDFAvail_IsFormAvail(*avail, hints);
+    if (nRet == PDF_FORM_ERROR || nRet == PDF_FORM_NOTAVAIL)
       return false;
-    }
-    int page_count = FPDF_GetPageCount(document_);
+
+    int page_count = FPDF_GetPageCount(*document);
     for (int i = 0; i < page_count; ++i) {
       nRet = PDF_DATA_NOTAVAIL;
-      while (nRet == PDF_DATA_NOTAVAIL) {
-        nRet = FPDFAvail_IsPageAvail(avail_, i, &hints_);
-      }
-      if (nRet == PDF_DATA_ERROR) {
+      while (nRet == PDF_DATA_NOTAVAIL)
+        nRet = FPDFAvail_IsPageAvail(*avail, i, hints);
+
+      if (nRet == PDF_DATA_ERROR)
         return false;
-      }
     }
   } else {
-    if (must_linearize) {
+    if (must_linearize)
       return false;
-    }
-    document_ = FPDF_LoadCustomDocument(&file_access_, password);
-    if (!document_) {
+
+    *document = FPDF_LoadCustomDocument(file_access, password);
+    if (!*document)
       return false;
-    }
   }
-  form_handle_ = SetupFormFillEnvironment(document_);
+  *form_handle = SetupFormFillEnvironment(*document);
 #ifdef PDF_ENABLE_XFA
   int docType = DOCTYPE_PDF;
-  if (FPDF_HasXFAField(document_, &docType)) {
+  if (FPDF_HasXFAField(*document, &docType)) {
     if (docType != DOCTYPE_PDF)
-      (void)FPDF_LoadXFA(document_);
+      (void)FPDF_LoadXFA(*document);
   }
 #endif  // PDF_ENABLE_XFA
-  (void)FPDF_GetDocPermissions(document_);
+  (void)FPDF_GetDocPermissions(*document);
   return true;
 }
 
@@ -286,16 +294,21 @@ void EmbedderTest::UnloadPage(FPDF_PAGE page) {
   page_reverse_map_.erase(it);
 }
 
-void EmbedderTest::TestSaved(int width, int height, const char* md5) {
+void EmbedderTest::TestSaved(int width,
+                             int height,
+                             const char* md5,
+                             const char* password) {
   FPDF_FILEACCESS file_access;
   memset(&file_access, 0, sizeof(file_access));
   file_access.m_FileLen = m_String.size();
   file_access.m_GetBlock = GetBlockFromString;
   file_access.m_Param = &m_String;
+  FX_FILEAVAIL file_avail;
+  FX_DOWNLOADHINTS hints;
 
-  m_SavedDocument = FPDF_LoadCustomDocument(&file_access, nullptr);
-  m_SavedForm = SetupFormFillEnvironment(m_SavedDocument);
-  ASSERT_TRUE(m_SavedDocument);
+  ASSERT_TRUE(OpenDocumentHelper(password, false, &file_avail, &hints,
+                                 &file_access, &m_SavedDocument, &m_SavedAvail,
+                                 &m_SavedForm));
   EXPECT_EQ(1, FPDF_GetPageCount(m_SavedDocument));
   m_SavedPage = FPDF_LoadPage(m_SavedDocument, 0);
   ASSERT_TRUE(m_SavedPage);
@@ -309,6 +322,7 @@ void EmbedderTest::CloseSaved() {
   FPDF_ClosePage(m_SavedPage);
   FPDFDOC_ExitFormFillEnvironment(m_SavedForm);
   FPDF_CloseDocument(m_SavedDocument);
+  FPDFAvail_Destroy(m_SavedAvail);
 }
 
 void EmbedderTest::TestAndCloseSaved(int width, int height, const char* md5) {
diff --git a/testing/embedder_test.h b/testing/embedder_test.h
index d2fd984a10..878e50bfdd 100644
--- a/testing/embedder_test.h
+++ b/testing/embedder_test.h
@@ -114,6 +114,15 @@ class EmbedderTest : public ::testing::Test,
   virtual void UnloadPage(FPDF_PAGE page);
 
  protected:
+  bool OpenDocumentHelper(const char* password,
+                          bool must_linearize,
+                          FX_FILEAVAIL* file_avail,
+                          FX_DOWNLOADHINTS* hints,
+                          FPDF_FILEACCESS* file_access,
+                          FPDF_DOCUMENT* document,
+                          FPDF_AVAIL* avail,
+                          FPDF_FORMHANDLE* form_handle);
+
   FPDF_FORMHANDLE SetupFormFillEnvironment(FPDF_DOCUMENT doc);
 
   // Return the hash of |bitmap|.
@@ -135,7 +144,10 @@ class EmbedderTest : public ::testing::Test,
                                 unsigned char* buf,
                                 unsigned long size);
 
-  void TestSaved(int width, int height, const char* md5);
+  void TestSaved(int width,
+                 int height,
+                 const char* md5,
+                 const char* password = nullptr);
   void CloseSaved();
   void TestAndCloseSaved(int width, int height, const char* md5);
 
@@ -159,6 +171,7 @@ class EmbedderTest : public ::testing::Test,
   FPDF_DOCUMENT m_SavedDocument;
   FPDF_PAGE m_SavedPage;
   FPDF_FORMHANDLE m_SavedForm;
+  FPDF_AVAIL m_SavedAvail;
 
  private:
   static void UnsupportedHandlerTrampoline(UNSUPPORT_INFO*, int type);