Get signals working in the EXTRA_DATA section of MHTML

content/browser/download/mhtml_generation_manager.cc

Index: content/browser/download/mhtml_generation_manager.cc
diff --git a/content/browser/download/mhtml_generation_manager.cc b/content/browser/download/mhtml_generation_manager.cc
index b989af11932a1e2c49ec1451e7d55560874b8f28..874f498b6c2e15f7d2044a98d33e695b6536c52e 100644
--- a/content/browser/download/mhtml_generation_manager.cc
+++ b/content/browser/download/mhtml_generation_manager.cc
@@ -21,10 +21,12 @@
 #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"
@@ -32,6 +34,14 @@
 #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: ";
+const char kContentTransferEncodingBinary[] =
+    "Content-Transfer-Encoding: binary";
+int kInvalidFileSize = -1;
+}  // namespace
+
 namespace content {
 
 // The class and all of its members live on the UI thread.  Only static methods
@@ -98,12 +108,27 @@ class MHTMLGenerationManager::Job : public RenderProcessHostObserver {
   // 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 MhtmlSaveStatus 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.
@@ -156,6 +181,9 @@ class MHTMLGenerationManager::Job : public RenderProcessHostObserver {
   // 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_;
@@ -185,6 +213,10 @@ MHTMLGenerationManager::Job::Job(int job_id,
   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() {
@@ -352,10 +384,11 @@ void MHTMLGenerationManager::Job::CloseFile(
   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);
 }
 
@@ -401,9 +434,11 @@ MhtmlSaveStatus MHTMLGenerationManager::Job::OnSerializeAsMHTMLResponse(
 
 // 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.
@@ -411,23 +446,91 @@ MHTMLGenerationManager::Job::CloseFileOnFileThread(MhtmlSaveStatus save_status,
     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;
+      return std::make_tuple(save_status, kInvalidFileSize);
+    }
+
+    // Write out the footer at the bottom of the file.
+    if (!WriteFooter(boundary, file)) {
+      save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+      return std::make_tuple(save_status, kInvalidFileSize);
+    }
   }
 
   // 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;
+  int64_t file_size = kInvalidFileSize;
+  MhtmlSaveStatus close_save_status = CloseFileIfValid(file, file_size);
+  if (save_status == MhtmlSaveStatus::SUCCESS &&
+      close_save_status == MhtmlSaveStatus::FILE_WRITTING_ERROR) {
+    save_status = close_save_status;
+  }
+
+  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
+    // content-transfer-encoding, 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%s\r\n", boundary.c_str(),
+        kContentLocation, part.content_location.c_str(), kContentType,
+        part.content_type.c_str(), kContentTransferEncodingBinary,

[Pete Williamson, 2017/04/01 01:02:58]

Thinking about this just a bit more, I think I'll move the
ContentTransferEncoding header into the string supplied by the client instead of
assuming it makes sense to set it to binary for all parts, new patch coming
Monday.

I'll also add a x-google-offline-pages-signal-collector header into the string
supplied by the prerendering offliner to make our parsing easier.

+        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
+MhtmlSaveStatus MHTMLGenerationManager::Job::CloseFileIfValid(
+    base::File& file,
+    int64_t& file_size) {
+  DCHECK_CURRENTLY_ON(BrowserThread::FILE);
   if (file.IsValid()) {
     file_size = file.GetLength();
     file.Close();
-  } else if (save_status == MhtmlSaveStatus::SUCCESS) {
-    save_status = MhtmlSaveStatus::FILE_CLOSING_ERROR;
+    return MhtmlSaveStatus::SUCCESS;
   }
 
-  return std::make_tuple(save_status, file_size);
+  return MhtmlSaveStatus::FILE_CLOSING_ERROR;
 }
 
 MHTMLGenerationManager* MHTMLGenerationManager::GetInstance() {