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..d614da87a2facd527f28578b4bde11e6cef90001 100644
--- a/content/browser/download/mhtml_generation_manager.cc
+++ b/content/browser/download/mhtml_generation_manager.cc
@@ -32,6 +32,45 @@
 #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;
+
+// Data fields used to build an additional MHTML part in the output file.
+struct MHTMLExtraPart {
+  std::string content_type;
+  std::string content_location;
+  std::string body;
+};
+
+// Class used as a data object for WebContents UserData to represent a MHTML
+// part that we plan to write into the output MHTML file.  Each MHTMLExtraPart
+// object in the contained vector lets us hold enough information to generate
+// one MHTML part.  Arbitrary extra MHTML parts to be added into the complete
+// file.  For instance, this might be used for gathering load time signals in
+// debug code for analysis.
+class MHTMLExtraData : public base::SupportsUserData::Data {
+ public:
+  MHTMLExtraData() {}
+  ~MHTMLExtraData() override {}
+
+  // Return the vector of parts to be disassembled.
+  std::vector<MHTMLExtraPart>& parts() { return parts_; }
+  const std::vector<MHTMLExtraPart>& parts() const { return parts_; }
+
+  // Get the data string out of the web contents. The web contents retains
+  // ownership of the vector.  Returns nullptr if none found.
+  static MHTMLExtraData* FromWebContents(content::WebContents* contents) {
+    return static_cast<MHTMLExtraData*>(
+        contents->GetUserData(&kMHTMLExtraDataKey));
+  }
+
+ private:
+  std::vector<MHTMLExtraPart> parts_;
+};
+}  // namespace
+
 namespace content {
 
 // The class and all of its members live on the UI thread.  Only static methods
@@ -101,9 +140,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 MHTMLExtraData* 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.
+  static bool WriteExtraDataPartsOnFileThread(const std::string& boundary,
+                                              base::File& file,
+                                              const MHTMLExtraData* 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 +202,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.
+  MHTMLExtraData* extra_data_;
+
   // RAII helper for registering this Job as a RenderProcessHost observer.
   ScopedObserver<RenderProcessHost, MHTMLGenerationManager::Job>
       observed_renderer_process_host_;
@@ -185,6 +234,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_data_ = MHTMLExtraData::FromWebContents(web_contents);
 }
 
 MHTMLGenerationManager::Job::~Job() {
@@ -355,7 +407,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,25 +453,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 MHTMLExtraData* extra_data) {
   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_data)) {
+      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 +493,69 @@ 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 MHTMLExtraData* extra_data) {
+  // Don't write an extra data part if there is none.
+  if (extra_data == nullptr)
+    return true;
+
+  const std::vector<MHTMLExtraPart>& extra_parts(extra_data->parts());
+  if (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();
 }
 
+void MHTMLGenerationManager::StashMHTMLPartForAdditionToPage(
+    content::WebContents* contents,
+    const std::string& content_type,
+    const std::string& content_location,
+    const std::string& body) {
+  MHTMLExtraPart part;
+  part.content_type = content_type;
+  part.content_location = content_location;
+  part.body = body;
+
+  MHTMLExtraData* extra_data = MHTMLExtraData::FromWebContents(contents);
+  // 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 here.
+    contents->SetUserData(&kMHTMLExtraDataKey,
+                          std::unique_ptr<MHTMLExtraData>(extra_data));
+  }
+  // Add this part to the list of parts to be saved out when the file is
+  // written.
+  extra_data->parts().push_back(part);
+}
+
 MHTMLGenerationManager::MHTMLGenerationManager() : next_job_id_(0) {}
 
 MHTMLGenerationManager::~MHTMLGenerationManager() {