Support partial loading in TestDocumentLoader

TestDocumentLoader currently loads the entire test PDF file as soon as
PDFiumEngine calls DocumentLoader::RequestData(0, 1). This change adds a
new TestDocumentLoader::SimulateLoadData(uint32_t) method that allows a
test to control how quickly the document loads.

A future test will take advantage of this to put the PDFiumEngine into a
partially-loaded state.

A test can request the TestDocumentLoader by calling PDFiumTestBase's
InitializeEngineWithoutLoading() method instead of InitializeEngine().

This change also removes the DocumentLoader::SetDocumentSize() method,
which only was used internally by DocumentLoaderImpl anyway.

Bug: 1051548
Change-Id: Ib5256aa70c468284e0f01964b409f2a85a9a066a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2067692
Commit-Queue: K Moon <kmoon@chromium.org>
Reviewed-by: Lei Zhang <thestig@chromium.org>
Reviewed-by: Henrique Nakashima <hnakashima@chromium.org>
Reviewed-by: Daniel Hosseinian <dhoss@chromium.org>
Cr-Commit-Position: refs/heads/master@{#751597}

pdf/document_loader.h

@@ -51,15 +51,14 @@ class DocumentLoader {
   virtual bool IsDataAvailable(uint32_t position, uint32_t size) const = 0;
 
   // Data request interface.
-  virtual void RequestData(uint32_t position, uint32_t size) {}
+  virtual void RequestData(uint32_t position, uint32_t size) = 0;
 
   virtual bool IsDocumentComplete() const = 0;
-  virtual void SetDocumentSize(uint32_t size) {}
   virtual uint32_t GetDocumentSize() const = 0;
   virtual uint32_t BytesReceived() const = 0;
 
   // Clear pending requests from the queue.
-  virtual void ClearPendingRequests() {}
+  virtual void ClearPendingRequests() = 0;
 };
 
 }  // namespace chrome_pdf

pdf/document_loader_impl.cc

@@ -98,14 +98,15 @@ bool DocumentLoaderImpl::Init(std::unique_ptr<URLLoaderWrapper> loader,
   loader_ = std::move(loader);
 
   if (!loader_->IsContentEncoded())
-    SetDocumentSize(std::max(0, loader_->GetContentLength()));
+    chunk_stream_.set_eof_pos(std::max(0, loader_->GetContentLength()));
 
   int64_t bytes_received = 0;
   int64_t total_bytes_to_be_received = 0;
   if (GetDocumentSize() == 0 &&
       loader_->GetDownloadProgress(&bytes_received,
                                    &total_bytes_to_be_received)) {
-    SetDocumentSize(std::max(0, static_cast<int>(total_bytes_to_be_received)));
+    chunk_stream_.set_eof_pos(
+        std::max(0, static_cast<int>(total_bytes_to_be_received)));
   }
 
   SetPartialLoadingEnabled(
@@ -122,10 +123,6 @@ bool DocumentLoaderImpl::IsDocumentComplete() const {
   return chunk_stream_.IsComplete();
 }
 
-void DocumentLoaderImpl::SetDocumentSize(uint32_t size) {
-  chunk_stream_.set_eof_pos(size);
-}
-
 uint32_t DocumentLoaderImpl::GetDocumentSize() const {
   return chunk_stream_.eof_pos();
 }
@@ -395,7 +392,7 @@ void DocumentLoaderImpl::ReadComplete() {
           chunk_stream_.filled_chunks().Last().end() * DataStream::kChunkSize,
           eof);
     }
-    SetDocumentSize(eof);
+    chunk_stream_.set_eof_pos(eof);
     if (eof == EndOfCurrentChunk())
       SaveChunkData();
   }

pdf/document_loader_impl.h

@@ -30,7 +30,6 @@ class DocumentLoaderImpl : public DocumentLoader {
   bool IsDataAvailable(uint32_t position, uint32_t size) const override;
   void RequestData(uint32_t position, uint32_t size) override;
   bool IsDocumentComplete() const override;
-  void SetDocumentSize(uint32_t size) override;
   uint32_t GetDocumentSize() const override;
   uint32_t BytesReceived() const override;
   void ClearPendingRequests() override;

pdf/pdfium/pdfium_engine.cc

@@ -121,9 +121,6 @@ constexpr base::TimeDelta kMaxProgressivePaintTime =
 constexpr base::TimeDelta kMaxInitialProgressivePaintTime =
     base::TimeDelta::FromMilliseconds(250);
 
-PDFiumEngine::CreateDocumentLoaderFunction
-    g_create_document_loader_for_testing = nullptr;
-
 template <class S>
 bool IsAboveOrDirectlyLeftOf(const S& lhs, const S& rhs) {
   return lhs.y() < rhs.y() || (lhs.y() == rhs.y() && lhs.x() < rhs.x());
@@ -432,10 +429,12 @@ PDFiumEngine::~PDFiumEngine() {
     FORM_DoDocumentAAction(form(), FPDFDOC_AACTION_WC);
 }
 
-// static
-void PDFiumEngine::SetCreateDocumentLoaderFunctionForTesting(
-    CreateDocumentLoaderFunction function) {
-  g_create_document_loader_for_testing = function;
+void PDFiumEngine::SetDocumentLoaderForTesting(
+    std::unique_ptr<DocumentLoader> loader) {
+  DCHECK(loader);
+  DCHECK(!doc_loader_);
+  doc_loader_ = std::move(loader);
+  doc_loader_set_for_testing_ = true;
 }
 
 bool PDFiumEngine::New(const char* url, const char* headers) {
@@ -577,9 +576,7 @@ bool PDFiumEngine::HandleDocumentLoad(const pp::URLLoader& loader) {
   password_tries_remaining_ = kMaxPasswordTries;
   process_when_pending_request_complete_ = true;
 
-  if (g_create_document_loader_for_testing) {
-    doc_loader_ = g_create_document_loader_for_testing(this);
-  } else {
+  if (!doc_loader_set_for_testing_) {
     auto loader_wrapper =
         std::make_unique<URLLoaderWrapperImpl>(GetPluginInstance(), loader);
     loader_wrapper->SetResponseHeaders(headers_);
@@ -1622,7 +1619,7 @@ void PDFiumEngine::StartFind(const std::string& text, bool case_sensitive) {
 
   // In unit tests, PPAPI is not initialized, so just call ContinueFind()
   // directly.
-  if (g_create_document_loader_for_testing) {
+  if (doc_loader_set_for_testing_) {
     ContinueFind(case_sensitive ? 1 : 0);
   } else {
     pp::CompletionCallback callback =

pdf/pdfium/pdfium_engine.h

@@ -55,10 +55,9 @@ class PDFiumEngine : public PDFEngine,
   PDFiumEngine(PDFEngine::Client* client, bool enable_javascript);
   ~PDFiumEngine() override;
 
-  using CreateDocumentLoaderFunction =
-      std::unique_ptr<DocumentLoader> (*)(DocumentLoader::Client* client);
-  static void SetCreateDocumentLoaderFunctionForTesting(
-      CreateDocumentLoaderFunction function);
+  // Replaces the normal DocumentLoader for testing. Must be called before
+  // HandleDocumentLoad().
+  void SetDocumentLoaderForTesting(std::unique_ptr<DocumentLoader> loader);
 
   // PDFEngine implementation.
   bool New(const char* url, const char* headers) override;
@@ -589,6 +588,7 @@ class PDFiumEngine : public PDFEngine,
   double current_zoom_ = 1.0;
 
   std::unique_ptr<DocumentLoader> doc_loader_;  // Main document's loader.
+  bool doc_loader_set_for_testing_ = false;
   std::string url_;
   std::string headers_;
   pp::CompletionCallbackFactory<PDFiumEngine> find_factory_;

pdf/pdfium/pdfium_test_base.cc

@@ -5,6 +5,7 @@
 #include "pdf/pdfium/pdfium_test_base.h"
 
 #include <memory>
+#include <utility>
 
 #include "build/build_config.h"
 #include "pdf/pdfium/pdfium_engine.h"
@@ -19,13 +20,6 @@ namespace chrome_pdf {
 
 namespace {
 
-const base::FilePath::CharType* g_test_pdf_name;
-
-std::unique_ptr<DocumentLoader> CreateTestDocumentLoader(
-    DocumentLoader::Client* client) {
-  return std::make_unique<TestDocumentLoader>(client, g_test_pdf_name);
-}
-
 bool IsValidLinkForTesting(const std::string& url) {
   return !url.empty();
 }
@@ -51,34 +45,47 @@ void PDFiumTestBase::SetUp() {
 }
 
 void PDFiumTestBase::TearDown() {
-  PDFiumEngine::SetCreateDocumentLoaderFunctionForTesting(nullptr);
   PDFiumPage::SetIsValidLinkFunctionForTesting(nullptr);
-  g_test_pdf_name = nullptr;
   FPDF_DestroyLibrary();
 }
 
 std::unique_ptr<PDFiumEngine> PDFiumTestBase::InitializeEngine(
     TestClient* client,
     const base::FilePath::CharType* pdf_name) {
-  SetDocumentForTest(pdf_name);
-  pp::URLLoader dummy_loader;
-  auto engine =
-      std::make_unique<PDFiumEngine>(client, /*enable_javascript=*/false);
-  client->set_engine(engine.get());
-  if (!engine->New("https://chromium.org/dummy.pdf", "") ||
-      !engine->HandleDocumentLoad(dummy_loader)) {
-    client->set_engine(nullptr);
-    return nullptr;
+  InitializeEngineResult result =
+      InitializeEngineWithoutLoading(client, pdf_name);
+  if (result.engine) {
+    // Incrementally read the PDF. To detect linearized PDFs, the first read
+    // should be at least 1024 bytes.
+    while (result.document_loader->SimulateLoadData(1024))
+      continue;
   }
-  return engine;
+  return std::move(result.engine);
 }
 
-void PDFiumTestBase::SetDocumentForTest(
+PDFiumTestBase::InitializeEngineResult
+PDFiumTestBase::InitializeEngineWithoutLoading(
+    TestClient* client,
     const base::FilePath::CharType* pdf_name) {
-  DCHECK(!g_test_pdf_name);
-  g_test_pdf_name = pdf_name;
-  PDFiumEngine::SetCreateDocumentLoaderFunctionForTesting(
-      &CreateTestDocumentLoader);
+  InitializeEngineResult result;
+
+  pp::URLLoader dummy_loader;
+  result.engine =
+      std::make_unique<PDFiumEngine>(client, /*enable_javascript=*/false);
+  client->set_engine(result.engine.get());
+
+  auto test_loader =
+      std::make_unique<TestDocumentLoader>(result.engine.get(), pdf_name);
+  result.document_loader = test_loader.get();
+  result.engine->SetDocumentLoaderForTesting(std::move(test_loader));
+
+  if (!result.engine->New("https://chromium.org/dummy.pdf", "") ||
+      !result.engine->HandleDocumentLoad(dummy_loader)) {
+    client->set_engine(nullptr);
+    result.engine = nullptr;
+    result.document_loader = nullptr;
+  }
+  return result;
 }
 
 void PDFiumTestBase::InitializePDFium() {
@@ -97,4 +104,14 @@ PDFiumPage* PDFiumTestBase::GetPDFiumPageForTest(PDFiumEngine* engine,
   return nullptr;
 }
 
+PDFiumTestBase::InitializeEngineResult::InitializeEngineResult() = default;
+
+PDFiumTestBase::InitializeEngineResult::InitializeEngineResult(
+    InitializeEngineResult&& other) = default;
+
+PDFiumTestBase::InitializeEngineResult& PDFiumTestBase::InitializeEngineResult::
+operator=(InitializeEngineResult&& other) = default;
+
+PDFiumTestBase::InitializeEngineResult::~InitializeEngineResult() = default;
+
 }  // namespace chrome_pdf

pdf/pdfium/pdfium_test_base.h

@@ -5,10 +5,11 @@
 #ifndef PDF_PDFIUM_PDFIUM_TEST_BASE_H_
 #define PDF_PDFIUM_PDFIUM_TEST_BASE_H_
 
+#include <stddef.h>
+
 #include <memory>
 
 #include "base/files/file_path.h"
-#include "base/macros.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace chrome_pdf {
@@ -16,37 +17,54 @@ namespace chrome_pdf {
 class PDFiumEngine;
 class PDFiumPage;
 class TestClient;
+class TestDocumentLoader;
 
 class PDFiumTestBase : public testing::Test {
  public:
   PDFiumTestBase();
+  PDFiumTestBase(const PDFiumTestBase&) = delete;
+  PDFiumTestBase& operator=(const PDFiumTestBase&) = delete;
   ~PDFiumTestBase() override;
 
   // Returns true when actually running in a Chrome OS environment.
   static bool IsRunningOnChromeOS();
 
  protected:
+  // Result of calling InitializeEngineWithoutLoading().
+  struct InitializeEngineResult {
+    InitializeEngineResult();
+    InitializeEngineResult(InitializeEngineResult&& other) noexcept;
+    InitializeEngineResult& operator=(InitializeEngineResult&& other) noexcept;
+    ~InitializeEngineResult();
+
+    // Initialized engine.
+    std::unique_ptr<PDFiumEngine> engine;
+
+    // Corresponding test document loader.
+    TestDocumentLoader* document_loader;
+  };
+
   // testing::Test:
   void SetUp() override;
   void TearDown() override;
 
   // Initializes a PDFiumEngine for use in testing with |client|. Loads a PDF
-  // named |pdf_name|.See TestDocumentLoader for more info about |pdf_name|.
+  // named |pdf_name|. See TestDocumentLoader for more info about |pdf_name|.
   std::unique_ptr<PDFiumEngine> InitializeEngine(
       TestClient* client,
       const base::FilePath::CharType* pdf_name);
 
+  // Initializes a PDFiumEngine as with InitializeEngine(), but defers loading
+  // until the test calls SimulateLoadData() on the returned TestDocumentLoader.
+  InitializeEngineResult InitializeEngineWithoutLoading(
+      TestClient* client,
+      const base::FilePath::CharType* pdf_name);
+
   // Returns the PDFiumPage for the page index
   PDFiumPage* GetPDFiumPageForTest(PDFiumEngine* engine, size_t page_index);
 
  private:
-  // Sets the PDF to load for a test. This must be called for tests that use
-  // TestDocumentLoader. See TestDocumentLoader for more info.
-  void SetDocumentForTest(const base::FilePath::CharType* pdf_name);
-
   void InitializePDFium();
-
-  DISALLOW_COPY_AND_ASSIGN(PDFiumTestBase);
 };
 
 }  // namespace chrome_pdf

pdf/test/test_document_loader.cc

@@ -4,11 +4,15 @@
 
 #include "pdf/test/test_document_loader.h"
 
+#include <stdint.h>
+
 #include "base/base_paths.h"
 #include "base/files/file_util.h"
 #include "base/logging.h"
 #include "base/path_service.h"
+#include "pdf/range_set.h"
 #include "pdf/url_loader_wrapper.h"
+#include "ui/gfx/range/range.h"
 
 namespace chrome_pdf {
 
@@ -27,9 +31,44 @@ TestDocumentLoader::TestDocumentLoader(
 
 TestDocumentLoader::~TestDocumentLoader() = default;
 
+// TODO(crbug.com/1056817): Consider faking out URLLoaderWrapper, to avoid
+// simulating the behavior of DocumentLoaderImpl (although that would result in
+// 64 KiB loads).
+bool TestDocumentLoader::SimulateLoadData(uint32_t max_bytes) {
+  CHECK_GT(max_bytes, 0U);
+  if (IsDocumentComplete())
+    return false;
+
+  // Approximate the behavior of DocumentLoaderImpl::ContinueDownload() by
+  // either reading from the start of the next pending range, or from the
+  // beginning of the document (skipping any received ranges).
+  RangeSet candidate_ranges(gfx::Range(
+      pending_ranges_.IsEmpty() ? 0 : pending_ranges_.First().start(),
+      GetDocumentSize()));
+  candidate_ranges.Subtract(received_ranges_);
+  CHECK(!candidate_ranges.IsEmpty());
+
+  gfx::Range request_range = candidate_ranges.First();
+  if (request_range.length() > max_bytes)
+    request_range.set_end(request_range.start() + max_bytes);
+
+  // Simulate fetching the requested range.
+  received_bytes_ += request_range.length();
+  received_ranges_.Union(request_range);
+  pending_ranges_.Subtract(request_range);
+  client_->OnNewDataReceived();
+
+  bool is_pending = !IsDocumentComplete();
+  if (is_pending)
+    client_->OnPendingRequestComplete();
+  else
+    client_->OnDocumentComplete();
+  return is_pending;
+}
+
 bool TestDocumentLoader::Init(std::unique_ptr<URLLoaderWrapper> loader,
                               const std::string& url) {
-  NOTREACHED();
+  NOTREACHED() << "PDFiumEngine skips this call when testing";
   return false;
 }
 
@@ -45,16 +84,25 @@ bool TestDocumentLoader::GetBlock(uint32_t position,
 
 bool TestDocumentLoader::IsDataAvailable(uint32_t position,
                                          uint32_t size) const {
-  return position < pdf_data_.size() && size <= pdf_data_.size() &&
-         position + size <= pdf_data_.size();
+  CHECK_LE(position, GetDocumentSize());
+  CHECK_LE(size, GetDocumentSize() - position);
+  gfx::Range range(position, position + size);
+  return range.is_empty() || received_ranges_.Contains(range);
 }
 
 void TestDocumentLoader::RequestData(uint32_t position, uint32_t size) {
-  client_->OnDocumentComplete();
+  if (IsDataAvailable(position, size))
+    return;
+
+  // DocumentLoaderImpl requests chunks of 64 KiB, but that is uninteresting for
+  // small test files, so use byte ranges instead.
+  RangeSet request_ranges(gfx::Range(position, position + size));
+  request_ranges.Subtract(received_ranges_);
+  pending_ranges_.Union(request_ranges);
 }
 
 bool TestDocumentLoader::IsDocumentComplete() const {
-  return true;
+  return BytesReceived() == GetDocumentSize();
 }
 
 uint32_t TestDocumentLoader::GetDocumentSize() const {
@@ -62,7 +110,11 @@ uint32_t TestDocumentLoader::GetDocumentSize() const {
 }
 
 uint32_t TestDocumentLoader::BytesReceived() const {
-  return pdf_data_.size();
+  return received_bytes_;
+}
+
+void TestDocumentLoader::ClearPendingRequests() {
+  pending_ranges_.Clear();
 }
 
 }  // namespace chrome_pdf

pdf/test/test_document_loader.h

@@ -5,11 +5,14 @@
 #ifndef PDF_TEST_TEST_DOCUMENT_LOADER_H_
 #define PDF_TEST_TEST_DOCUMENT_LOADER_H_
 
+#include <stdint.h>
+
 #include <memory>
 #include <string>
 
 #include "base/files/file_path.h"
 #include "pdf/document_loader.h"
+#include "pdf/range_set.h"
 
 namespace chrome_pdf {
 
@@ -21,6 +24,10 @@ class TestDocumentLoader : public DocumentLoader {
                      const base::FilePath::StringType& pdf_name);
   ~TestDocumentLoader() override;
 
+  // Simulates loading up to |max_bytes| more data, returning `true` if there is
+  // more data to load (that is, IsDocumentComplete() returns `false`).
+  bool SimulateLoadData(uint32_t max_bytes);
+
   // DocumentLoader:
   bool Init(std::unique_ptr<URLLoaderWrapper> loader,
             const std::string& url) override;
@@ -30,10 +37,16 @@ class TestDocumentLoader : public DocumentLoader {
   bool IsDocumentComplete() const override;
   uint32_t GetDocumentSize() const override;
   uint32_t BytesReceived() const override;
+  void ClearPendingRequests() override;
 
  private:
   Client* const client_;
   std::string pdf_data_;
+
+  // Not using ChunkStream, for more fine-grained control over request size.
+  uint32_t received_bytes_ = 0;
+  RangeSet received_ranges_;
+  RangeSet pending_ranges_;
 };
 
 }  // namespace chrome_pdf