Get signals working in the EXTRA_DATA section of MHTML

content/browser/download/mhtml_generation_manager.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "content/browser/download/mhtml_generation_manager.h"
 
 #include <map>
 #include <queue>
 #include <utility>
 
 #include "base/bind.h"
 #include "base/files/file.h"
 #include "base/guid.h"
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/scoped_observer.h"
 #include "base/stl_util.h"
 #include "base/strings/string_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/time/time.h"
 #include "base/trace_event/trace_event.h"
 #include "content/browser/bad_message.h"
+#include "content/browser/download/mhtml_extra_parts_impl.h"
 #include "content/browser/frame_host/frame_tree_node.h"
 #include "content/browser/frame_host/render_frame_host_impl.h"
 #include "content/common/frame_messages.h"
 #include "content/public/browser/browser_thread.h"
+#include "content/public/browser/mhtml_extra_parts.h"
 #include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/render_process_host.h"
 #include "content/public/browser/render_process_host_observer.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/common/mhtml_generation_params.h"
 #include "net/base/mime_util.h"
 
+namespace {
+const char kContentLocation[] = "Content-Location: ";
+const char kContentType[] = "Content-Type: ";
+int kInvalidFileSize = -1;
+}  // namespace
+
 namespace content {
 
 // The class and all of its members live on the UI thread.  Only static methods
 // are executed on other threads.
 class MHTMLGenerationManager::Job : public RenderProcessHostObserver {
  public:
   Job(int job_id,
       WebContents* web_contents,
       const MHTMLGenerationParams& params,
       const GenerateMHTMLCallback& callback);
@@ skipping 46 matching lines @@
   void MarkAsFinished();
 
   void ReportRendererMainThreadTime(base::TimeDelta renderer_main_thread_time);
 
  private:
   // Writes the MHTML footer to the file and closes it.
   //
   // Note: The same |boundary| marker must be used for all "boundaries" -- in
   // the header, parts and footer -- that belong to the same MHTML document (see
   // also rfc1341, section 7.2.1, "boundary" description).
-  static std::tuple<MhtmlSaveStatus, int64_t> CloseFileOnFileThread(
+  static std::tuple<MhtmlSaveStatus, int64_t> FinalizeAndCloseFileOnFileThread(
       MhtmlSaveStatus save_status,
       const std::string& boundary,
-      base::File file);
+      base::File file,
+      const MHTMLExtraPartsImpl* extra_parts);
   void AddFrame(RenderFrameHost* render_frame_host);
 
+  // If we have any extra MHTML parts to write out, write them into the file
+  // while on the file thread.  Returns true for success, or if there is no data
+  // to write.
+  static bool WriteExtraDataParts(const std::string& boundary,
+                                  base::File& file,
+                                  const MHTMLExtraPartsImpl* extra_parts);
+
+  // Writes the footer into the MHTML file.  Returns false for faiulre.
+  static bool WriteFooter(const std::string& boundary, base::File& file);
+
+  // Close the MHTML file if it looks good, setting the size param.  Returns
+  // false for failure.
+  static bool CloseFileIfValid(base::File& file, int64_t* file_size);
+
   // Creates a new map with values (content ids) the same as in
   // |frame_tree_node_to_content_id_| map, but with the keys translated from
   // frame_tree_node_id into a |site_instance|-specific routing_id.
   std::map<int, std::string> CreateFrameRoutingIdToContentId(
       SiteInstance* site_instance);
 
   // Id used to map renderer responses to jobs.
   // See also MHTMLGenerationManager::id_to_job_ map.
   const int job_id_;
 
@@ skipping 32 matching lines @@
   std::string salt_;
 
   // The callback to call once generation is complete.
   const GenerateMHTMLCallback callback_;
 
   // Whether the job is finished (set to true only for the short duration of
   // time between MHTMLGenerationManager::JobFinished is called and the job is
   // destroyed by MHTMLGenerationManager::OnFileClosed).
   bool is_finished_;
 
+  // Any extra data parts that should be emitted into the output MHTML.
+  MHTMLExtraPartsImpl* extra_parts_;
+
   // RAII helper for registering this Job as a RenderProcessHost observer.
   ScopedObserver<RenderProcessHost, MHTMLGenerationManager::Job>
       observed_renderer_process_host_;
 
   DISALLOW_COPY_AND_ASSIGN(Job);
 };
 
 MHTMLGenerationManager::Job::Job(int job_id,
                                  WebContents* web_contents,
                                  const MHTMLGenerationParams& params,
                                  const GenerateMHTMLCallback& callback)
     : job_id_(job_id),
       creation_time_(base::TimeTicks::Now()),
       params_(params),
       frame_tree_node_id_of_busy_frame_(FrameTreeNode::kFrameTreeNodeInvalidId),
       mhtml_boundary_marker_(net::GenerateMimeMultipartBoundary()),
       salt_(base::GenerateGUID()),
       callback_(callback),
       is_finished_(false),
       observed_renderer_process_host_(this) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   web_contents->ForEachFrame(base::Bind(
       &MHTMLGenerationManager::Job::AddFrame,
       base::Unretained(this)));  // Safe because ForEachFrame is synchronous.
 
   // Main frame needs to be processed first.
   DCHECK(!pending_frame_tree_node_ids_.empty());
   DCHECK(FrameTreeNode::GloballyFindByID(pending_frame_tree_node_ids_.front())
              ->parent() == nullptr);
+
+  // Save off any extra data.
+  extra_parts_ = static_cast<MHTMLExtraPartsImpl*>(
+      MHTMLExtraParts::FromWebContents(web_contents));
 }
 
 MHTMLGenerationManager::Job::~Job() {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
 }
 
 std::map<int, std::string>
 MHTMLGenerationManager::Job::CreateFrameRoutingIdToContentId(
     SiteInstance* site_instance) {
   std::map<int, std::string> result;
@@ skipping 147 matching lines @@
     if (save_status == MhtmlSaveStatus::SUCCESS)
       save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
     callback.Run(std::make_tuple(save_status, -1));
     return;
   }
 
   // If no previous error occurred the boundary should be sent.
   BrowserThread::PostTaskAndReplyWithResult(
       BrowserThread::FILE, FROM_HERE,
       base::Bind(
-          &MHTMLGenerationManager::Job::CloseFileOnFileThread, save_status,
+          &MHTMLGenerationManager::Job::FinalizeAndCloseFileOnFileThread,
+          save_status,
           (save_status == MhtmlSaveStatus::SUCCESS ? mhtml_boundary_marker_
                                                    : std::string()),
-          base::Passed(&browser_file_)),
+          base::Passed(&browser_file_), extra_parts_),
       callback);
 }
 
 bool MHTMLGenerationManager::Job::IsMessageFromFrameExpected(
     RenderFrameHostImpl* sender) {
   int sender_id = sender->frame_tree_node()->frame_tree_node_id();
   if (sender_id != frame_tree_node_id_of_busy_frame_)
     return false;
 
   // We only expect one message per frame - let's make sure subsequent messages
@@ skipping 25 matching lines @@
 
   // Report success if all frames have been processed.
   if (pending_frame_tree_node_ids_.empty())
     return MhtmlSaveStatus::SUCCESS;
 
   return SendToNextRenderFrame();
 }
 
 // static
 std::tuple<MhtmlSaveStatus, int64_t>
-MHTMLGenerationManager::Job::CloseFileOnFileThread(MhtmlSaveStatus save_status,
-                                                   const std::string& boundary,
-                                                   base::File file) {
+MHTMLGenerationManager::Job::FinalizeAndCloseFileOnFileThread(
+    MhtmlSaveStatus save_status,
+    const std::string& boundary,
+    base::File file,
+    const MHTMLExtraPartsImpl* extra_parts) {
   DCHECK_CURRENTLY_ON(BrowserThread::FILE);
 
   // If no previous error occurred the boundary should have been provided.
   if (save_status == MhtmlSaveStatus::SUCCESS) {
     TRACE_EVENT0("page-serialization",
                  "MHTMLGenerationManager::Job MHTML footer writing");
     DCHECK(!boundary.empty());
-    std::string footer = base::StringPrintf("--%s--\r\n", boundary.c_str());
-    DCHECK(base::IsStringASCII(footer));
-    if (file.WriteAtCurrentPos(footer.data(), footer.size()) < 0)
+
+    // Write the extra data into a part of its own, if we have any.
+    if (!WriteExtraDataParts(boundary, file, extra_parts)) {

[fgorski, 2017/04/05 22:42:38]

minor thing which I didn't realize until now that I see it. This could pretty
much be a single if statement with both calls:
if (!WriteExtraDataParts(boundary, file, extra_parts) ||
    !WriteFooter(boundary, file)) {
  save_status = ...;
}

Comments can be put on top, but the names of methods are good in this case, and
comments don't really add stuff here.

Consider applying if you are comfortable with such change.

       save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+    }
+
+    // Write out the footer at the bottom of the file.
+    if (save_status == MhtmlSaveStatus::SUCCESS &&
+        !WriteFooter(boundary, file)) {
+      save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+    }
   }
 
   // If the file is still valid try to close it. Only update the status if that
   // won't hide an earlier error.
-  int64_t file_size = -1;
-  if (file.IsValid()) {
-    file_size = file.GetLength();
-    file.Close();
-  } else if (save_status == MhtmlSaveStatus::SUCCESS) {
+  int64_t file_size = kInvalidFileSize;
+  if (!CloseFileIfValid(file, &file_size) &&
+      save_status == MhtmlSaveStatus::SUCCESS) {
     save_status = MhtmlSaveStatus::FILE_CLOSING_ERROR;
   }
 
   return std::make_tuple(save_status, file_size);
 }
 
+// static
+bool MHTMLGenerationManager::Job::WriteExtraDataParts(
+    const std::string& boundary,
+    base::File& file,
+    const MHTMLExtraPartsImpl* extra_parts) {
+  DCHECK_CURRENTLY_ON(BrowserThread::FILE);
+  // Don't write an extra data part if there is none.
+  if (extra_parts == nullptr)
+    return true;
+
+  const std::vector<MHTMLExtraDataPart>& extra_data_parts(extra_parts->parts());
+  if (extra_data_parts.empty())
+    return true;
+
+  std::string serialized_extra_data_parts;
+
+  // For each extra part, serialize that part and add to our accumulator
+  // string.
+  for (auto part : extra_data_parts) {
+    // Write a newline, then a boundary, another newline, then the content
+    // location, another newline, the content type, another newline, the another
+    // newline, the extra data string, and end with a newline.
+    std::string serialized_extra_data_part = base::StringPrintf(
+        "--%s\r\n%s%s\r\n%s%s\r\n%s\r\n", boundary.c_str(), kContentLocation,
+        part.content_location.c_str(), kContentType, part.content_type.c_str(),
+        part.body.c_str());
+    DCHECK(base::IsStringASCII(serialized_extra_data_part));
+
+    serialized_extra_data_parts += serialized_extra_data_part;
+  }
+
+  // Write the string into the file.  Returns false if we failed the write.
+  return (file.WriteAtCurrentPos(serialized_extra_data_parts.data(),
+                                 serialized_extra_data_parts.size()) >= 0);
+}
+
+// static
+bool MHTMLGenerationManager::Job::WriteFooter(const std::string& boundary,
+                                              base::File& file) {
+  DCHECK_CURRENTLY_ON(BrowserThread::FILE);
+  std::string footer = base::StringPrintf("--%s--\r\n", boundary.c_str());
+  DCHECK(base::IsStringASCII(footer));
+  return (file.WriteAtCurrentPos(footer.data(), footer.size()) >= 0);
+}
+
+// static
+bool MHTMLGenerationManager::Job::CloseFileIfValid(base::File& file,
+                                                   int64_t* file_size) {
+  DCHECK_CURRENTLY_ON(BrowserThread::FILE);
+  DCHECK(file_size);
+  if (file.IsValid()) {
+    *file_size = file.GetLength();
+    file.Close();
+    return true;
+  }
+
+  return false;
+}
+
 MHTMLGenerationManager* MHTMLGenerationManager::GetInstance() {
   return base::Singleton<MHTMLGenerationManager>::get();
 }
 
 MHTMLGenerationManager::MHTMLGenerationManager() : next_job_id_(0) {}
 
 MHTMLGenerationManager::~MHTMLGenerationManager() {
 }
 
 void MHTMLGenerationManager::SaveMHTML(WebContents* web_contents,
@@ skipping 149 matching lines @@
   return iter->second.get();
 }
 
 void MHTMLGenerationManager::RenderProcessExited(Job* job) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   DCHECK(job);
   JobFinished(job, MhtmlSaveStatus::RENDER_PROCESS_EXITED);
 }
 
 }  // namespace content