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>
 
@@ skipping 14 matching lines @@
 #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/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: ";
+const int kMHTMLExtraDataKey = 0;
+}  // 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 49 matching lines @@
 
  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(
       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.
   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.
+  std::vector<MHTMLExtraPart>* 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_ = MHTMLExtraData::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 150 matching lines @@
     return;
   }
 
   // If no previous error occurred the boundary should be sent.
   BrowserThread::PostTaskAndReplyWithResult(
       BrowserThread::FILE, FROM_HERE,
       base::Bind(
           &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);
 }
 
 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::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();
   } else if (save_status == MhtmlSaveStatus::SUCCESS) {
     save_status = MhtmlSaveStatus::FILE_CLOSING_ERROR;
   }
 
   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() {
 }
 
 void MHTMLGenerationManager::SaveMHTML(WebContents* web_contents,
                                        const MHTMLGenerationParams& params,
                                        const GenerateMHTMLCallback& callback) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
 
@@ skipping 145 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