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..c19599b51f61828c6b32e390ea9d2dbfd5c38cb7 100644
--- a/content/browser/download/mhtml_generation_manager.cc
+++ b/content/browser/download/mhtml_generation_manager.cc
@@ -21,6 +21,7 @@
 #include "base/time/time.h"
 #include "base/trace_event/trace_event.h"
 #include "content/browser/bad_message.h"
+#include "content/browser/download/mhtml_extra_data.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"
@@ -32,6 +33,10 @@
 #include "content/public/common/mhtml_generation_params.h"
 #include "net/base/mime_util.h"
 
+namespace {
+const char kContentLocation[] = "Content-Location: ";
+}  // namespace
+
 namespace content {
 
 // The class and all of its members live on the UI thread.  Only static methods
@@ -101,9 +106,16 @@ class MHTMLGenerationManager::Job : public RenderProcessHostObserver {
   static std::tuple<MhtmlSaveStatus, int64_t> CloseFileOnFileThread(
       MhtmlSaveStatus save_status,
       const std::string& boundary,
-      base::File file);
+      base::File file,
+      const std::string& extra_data);
   void AddFrame(RenderFrameHost* render_frame_host);
 
+  // If we have any extra data sections to write out, write them into the file
+  // while on the file thread.  Returns true for success.
+  static bool WriteExtraDataSectionOnFileThread(const std::string& boundary,
+                                                base::File& file,
+                                                const std::string& extra_data);
+
   // 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 +168,9 @@ class MHTMLGenerationManager::Job : public RenderProcessHostObserver {
   // destroyed by MHTMLGenerationManager::OnFileClosed).
   bool is_finished_;
 
+  // Any extra data section that should be emitted into the output MHTML.
+  std::string extra_data_;
+
   // RAII helper for registering this Job as a RenderProcessHost observer.
   ScopedObserver<RenderProcessHost, MHTMLGenerationManager::Job>
       observed_renderer_process_host_;
@@ -185,6 +200,14 @@ 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
+  base::SupportsUserData::Data* found_user_data =
+      web_contents->GetUserData(content::kMHTMLExtraDataKey);
+  MHTMLExtraData* found_data = static_cast<MHTMLExtraData*>(found_user_data);
+
+  if (found_data != nullptr)
+    extra_data_ = found_data->data_string;
 }
 
 MHTMLGenerationManager::Job::~Job() {
@@ -355,7 +378,7 @@ void MHTMLGenerationManager::Job::CloseFile(
           &MHTMLGenerationManager::Job::CloseFileOnFileThread, save_status,
           (save_status == MhtmlSaveStatus::SUCCESS ? mhtml_boundary_marker_
                                                    : std::string()),
-          base::Passed(&browser_file_)),
+          base::Passed(&browser_file_), extra_data_),
       callback);
 }
 
@@ -401,9 +424,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::CloseFileOnFileThread(
+    MhtmlSaveStatus save_status,
+    const std::string& boundary,
+    base::File file,
+    const std::string& extra_data) {
   DCHECK_CURRENTLY_ON(BrowserThread::FILE);
 
   // If no previous error occurred the boundary should have been provided.
@@ -411,6 +436,11 @@ MHTMLGenerationManager::Job::CloseFileOnFileThread(MhtmlSaveStatus save_status,
     TRACE_EVENT0("page-serialization",
                  "MHTMLGenerationManager::Job MHTML footer writing");
     DCHECK(!boundary.empty());
+
+    // Write the extra data into a section of its own, if we have any.
+    if (!WriteExtraDataSectionOnFileThread(boundary, file, extra_data))
+      save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+
     std::string footer = base::StringPrintf("--%s--\r\n", boundary.c_str());
     DCHECK(base::IsStringASCII(footer));
     if (file.WriteAtCurrentPos(footer.data(), footer.size()) < 0)
@@ -430,6 +460,32 @@ MHTMLGenerationManager::Job::CloseFileOnFileThread(MhtmlSaveStatus save_status,
   return std::make_tuple(save_status, file_size);
 }
 
+// static
+bool MHTMLGenerationManager::Job::WriteExtraDataSectionOnFileThread(
+    const std::string& boundary,
+    base::File& file,
+    const std::string& extra_data) {
+  // Don't write an extra data section if there is none.
+  if (extra_data.empty())
+    return true;
+
+  // TODO: What value do I use for content location?
+  std::string content_location_value;
+
+  // Write a newline, then a boundary, then the content location, another
+  // newline, the extra data string, and end with a newline.
+  std::string extra_data_section = base::StringPrintf(
+      "--%s--\r\n%s%s\r\n%s\r\n", boundary.c_str(), kContentLocation,
+      content_location_value.c_str(), extra_data.c_str());
+  // Now actually write the string into the file.
+  DCHECK(base::IsStringASCII(extra_data_section));
+  if (file.WriteAtCurrentPos(extra_data_section.data(),
+                             extra_data_section.size()) < 0)
+    return false;
+
+  return true;
+}
+
 MHTMLGenerationManager* MHTMLGenerationManager::GetInstance() {
   return base::Singleton<MHTMLGenerationManager>::get();
 }