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..5fdd801bb8768af79063cecb38f476f7ce7a0ec6 100644
--- a/content/browser/download/mhtml_generation_manager.cc
+++ b/content/browser/download/mhtml_generation_manager.cc
@@ -32,6 +32,12 @@
 #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 int kMHTMLExtraDataKey = 0;
+}  // namespace
+
 namespace content {
 
 // The class and all of its members live on the UI thread.  Only static methods
@@ -101,9 +107,17 @@ 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::vector<MHTMLGenerationManager::MHTMLExtraPart>* extra_parts);
   void AddFrame(RenderFrameHost* render_frame_host);
 
+  // If we have any extra data parts to write out, write them into the file
+  // while on the file thread.  Returns true for success.
+  static bool WriteExtraDataPartsOnFileThread(
+      const std::string& boundary,
+      base::File& file,
+      const std::vector<MHTMLExtraPart>* extra_parts);

[fgorski, 2017/03/20 20:43:24]

Why is this vector passed by pointer?

[Pete Williamson, 2017/03/20 21:43:52]

When we get the user data from the web contents, that comes by pointer.  It
seems more parallel with the surrounding code to pass this by pointer too.  Do
you have a preference for a reference?

Also, it passes cleanly as null if the data is not present (which is probably
why the user data underneath used a pointer).

+
   // 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 +170,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.
+  std::vector<MHTMLExtraPart>* extra_parts_;
+
   // RAII helper for registering this Job as a RenderProcessHost observer.
   ScopedObserver<RenderProcessHost, MHTMLGenerationManager::Job>
       observed_renderer_process_host_;
@@ -185,6 +202,9 @@ 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_ = MHTMLExtraData::FromWebContents(web_contents);
 }
 
 MHTMLGenerationManager::Job::~Job() {
@@ -355,7 +375,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_parts_),
       callback);
 }
 
@@ -401,25 +421,36 @@ 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::vector<MHTMLExtraPart>* extra_parts) {
   DCHECK_CURRENTLY_ON(BrowserThread::FILE);
+  int64_t file_size = -1;
 
   // 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());
+
+    // Write the extra data into a part of its own, if we have any.
+    if (!WriteExtraDataPartsOnFileThread(boundary, file, extra_parts)) {
+      save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+      return std::make_tuple(save_status, file_size);
+    }
+
     std::string footer = base::StringPrintf("--%s--\r\n", boundary.c_str());
     DCHECK(base::IsStringASCII(footer));
-    if (file.WriteAtCurrentPos(footer.data(), footer.size()) < 0)
+    if (file.WriteAtCurrentPos(footer.data(), footer.size()) < 0) {
       save_status = MhtmlSaveStatus::FILE_WRITTING_ERROR;
+      return std::make_tuple(save_status, file_size);
+    }
   }
 
   // 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();
@@ -430,10 +461,72 @@ MHTMLGenerationManager::Job::CloseFileOnFileThread(MhtmlSaveStatus save_status,
   return std::make_tuple(save_status, file_size);
 }
 
+// static
+bool MHTMLGenerationManager::Job::WriteExtraDataPartsOnFileThread(
+    const std::string& boundary,
+    base::File& file,
+    const std::vector<MHTMLExtraPart>* extra_parts) {
+  // Don't write an extra data part if there is none.
+  if (extra_parts == nullptr || extra_parts->empty())
+    return true;
+
+  std::string extra_parts_serialized;
+
+  // For each extra part, serialize that part and add to our accumulator
+  // string.
+  for (auto part : *extra_parts) {
+    // Write a newline, then a boundary, another newline, then the content
+    // location, another newline, the content type, another newline, the extra
+    // data string, and end with a newline.
+    std::string 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(extra_data_part));
+
+    extra_parts_serialized += extra_data_part;
+  }
+
+  // Write the string into the file.  Returns false if we failed the write.
+  return (file.WriteAtCurrentPos(extra_parts_serialized.data(),
+                                 extra_parts_serialized.size()) >= 0);
+}
+
 MHTMLGenerationManager* MHTMLGenerationManager::GetInstance() {
   return base::Singleton<MHTMLGenerationManager>::get();
 }
 
+MHTMLGenerationManager::MHTMLExtraData::MHTMLExtraData() {}
+
+MHTMLGenerationManager::MHTMLExtraData::~MHTMLExtraData() {}
+
+std::vector<MHTMLGenerationManager::MHTMLExtraPart>*
+MHTMLGenerationManager::MHTMLExtraData::FromWebContents(

[fgorski, 2017/03/20 20:43:24]

Why not return a pointer to extra data?

[Pete Williamson, 2017/03/20 21:43:52]

What the client of this method needs is not an ExtraData, but the underlying
vector, so why not make life easier for the caller, and encapsulate out the
ExtraData here?

+    content::WebContents* contents) {
+  MHTMLExtraData* extra_data =
+      static_cast<MHTMLExtraData*>(contents->GetUserData(&kMHTMLExtraDataKey));
+  if (extra_data == nullptr)
+    return nullptr;
+  return &(extra_data->parts_);
+}
+
+void MHTMLGenerationManager::MHTMLExtraData::AddPartToWebContents(
+    content::WebContents* contents,
+    MHTMLGenerationManager::MHTMLExtraPart& part) {
+  MHTMLExtraData* extra_data =
+      static_cast<MHTMLExtraData*>(contents->GetUserData(&kMHTMLExtraDataKey));
+  // If we don't already have any extra data, add it now.
+  if (extra_data == nullptr) {
+    extra_data = new MHTMLExtraData();
+    // Note that we hold onto extra_data until the end of this function, even
+    // though we transfer it to the web contents user data.
+    contents->SetUserData(&kMHTMLExtraDataKey,
+                          std::unique_ptr<MHTMLExtraData>(extra_data));
+  }
+  // Add this part to the list of parts to be saved out.
+  extra_data->parts_.push_back(part);
+}
+
 MHTMLGenerationManager::MHTMLGenerationManager() : next_job_id_(0) {}
 
 MHTMLGenerationManager::~MHTMLGenerationManager() {