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)) {
       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