Content API changes to improve DOM stitching in ThreatDetails code.

chrome/browser/safe_browsing/threat_details.cc

Index: chrome/browser/safe_browsing/threat_details.cc
diff --git a/chrome/browser/safe_browsing/threat_details.cc b/chrome/browser/safe_browsing/threat_details.cc
index 891651613248dd4c6fbdeb9902b93093edb83cfc..a40d5f21afaf13008557f42bd21d006ac5e04ee3 100644
--- a/chrome/browser/safe_browsing/threat_details.cc
+++ b/chrome/browser/safe_browsing/threat_details.cc
@@ -22,6 +22,7 @@
 #include "content/public/browser/navigation_controller.h"
 #include "content/public/browser/navigation_entry.h"
 #include "content/public/browser/render_frame_host.h"
+#include "content/public/browser/render_process_host.h"
 #include "content/public/browser/web_contents.h"
 #include "net/url_request/url_request_context_getter.h"
 
@@ -185,7 +186,6 @@ ThreatDetails::ThreatDetails(
       cache_result_(false),
       did_proceed_(false),
       num_visits_(0),
-      ambiguous_dom_(false),
       cache_collector_(new ThreatDetailsCacheCollector) {
   redirects_collector_ = new ThreatDetailsRedirectsCollector(
       history_service ? history_service->AsWeakPtr()
@@ -286,8 +286,9 @@ ClientSafeBrowsingReportRequest::Resource* ThreatDetails::AddUrl(
 }
 
 void ThreatDetails::AddDomElement(
+    const int process_id,
     const int frame_tree_node_id,
-    const std::string& frame_url,
+    const int other_frame_routing_id,
     const int element_node_id,
     const std::string& tagname,
     const int parent_element_node_id,
@@ -314,12 +315,15 @@ void ThreatDetails::AddDomElement(
   if (resource) {
     cur_element->set_resource_id(resource->id());
 
-    // For iframes, remember that this HTML Element represents an iframe with a
-    // specific URL. Elements from a frame with this URL are children of this
-    // element.
-    if (is_frame &&
-        !base::ContainsKey(iframe_src_to_element_map_, resource->url())) {
-      iframe_src_to_element_map_[resource->url()] = cur_element;
+    // For iframes, lookup the frame tree node id of the render frame that
+    // handled that iframe's content. This must be done on the UI thread, and
+    // will update a map of |element_key| to |frame_tree_node_id|. A second pass
+    // is done to update the |elements_| list using this mapping.
+    if (is_frame) {
+      BrowserThread::PostTask(
+          BrowserThread::UI, FROM_HERE,
+          base::Bind(&ThreatDetails::LookupOtherFrameId, this, element_key,
+                     process_id, other_frame_routing_id));
     }
   }
 
@@ -328,23 +332,11 @@ void ThreatDetails::AddDomElement(
   HTMLElement* parent_element = nullptr;
   if (parent_element_node_id == 0) {
     // No parent indicates that this element is at the top of the current frame.
-    // This frame could be a child of an iframe in another frame, or it could be
-    // at the root of the whole page. If we have a frame URL then we can try to
-    // map this element to its parent.
-    if (!frame_url.empty()) {
-      // First, remember that this element is at the top-level of a frame with
-      // our frame URL.
-      document_url_to_children_map_[frame_url].insert(cur_element->id());
-
-      // Now check if the frame URL matches the src URL of an iframe elsewhere.
-      // This means that we processed the parent iframe element earlier, so we
-      // can add ourselves as a child of that iframe.
-      // If no such iframe exists, it could be processed later, or this element
-      // is in the top-level frame and truly has no parent.
-      if (base::ContainsKey(iframe_src_to_element_map_, frame_url)) {
-        parent_element = iframe_src_to_element_map_[frame_url];
-      }
-    }
+    // Remember that this is a top-level element of the frame with the
+    // current |frame_tree_node_id|. If this element is inside an iframe, a
+    // second pass will insert this element as a child of its parent iframe.
+    frame_tree_id_to_children_map_[frame_tree_node_id].insert(
+        cur_element->id());
   } else {
     // We have a parent ID, so this element is just a child of something inside
     // of our current frame. We can easily lookup our parent.
@@ -369,20 +361,6 @@ void ThreatDetails::AddDomElement(
       parent_element->add_child_ids(cur_element->id());
     }
   }
-
-  // Finally, we need to check if the current element is the parent of some
-  // other elements that came in from another frame earlier. This only happens
-  // if we are an iframe, and our src URL exists in
-  // document_url_to_children_map_. If there is a match, then all of the
-  // children in that map belong to us.
-  if (is_frame && resource &&
-      base::ContainsKey(document_url_to_children_map_, resource->url())) {
-    const std::unordered_set<int>& child_ids =
-        document_url_to_children_map_[resource->url()];
-    for (const int child_id : child_ids) {
-      cur_element->add_child_ids(child_id);
-    }
-  }
 }
 
 void ThreatDetails::StartCollection() {
@@ -456,12 +434,13 @@ void ThreatDetails::OnReceivedThreatDOMDetails(
   // of our data structures (eg GetSerializedReport).
   BrowserThread::PostTask(
       BrowserThread::IO, FROM_HERE,
-      base::BindOnce(&ThreatDetails::AddDOMDetails, this,
-                     sender->GetFrameTreeNodeId(),
-                     sender->GetLastCommittedURL(), params));
+      base::Bind(&ThreatDetails::AddDOMDetails, this,
+                 sender->GetProcess()->GetID(), sender->GetFrameTreeNodeId(),
+                 sender->GetLastCommittedURL(), params));
 }
 
 void ThreatDetails::AddDOMDetails(
+    const int process_id,
     const int frame_tree_node_id,
     const GURL& frame_last_committed_url,
     const std::vector<SafeBrowsingHostMsg_ThreatDOMDetails_Node>& params) {
@@ -482,22 +461,6 @@ void ThreatDetails::AddDOMDetails(
   if (params.empty())
     return;
 
-  // Try to deduce the URL that the render frame was handling. First check if
-  // the summary node from the renderer has a document URL. If not, try looking
-  // at the last committed URL of the frame.
-  GURL frame_url;
-  if (IsReportableUrl(params.back().url)) {
-    frame_url = params.back().url;
-  } else if (IsReportableUrl(frame_last_committed_url)) {
-    frame_url = frame_last_committed_url;
-  }
-
-  // If we can't figure out which URL the frame was rendering then we don't know
-  // where these elements belong in the hierarchy. The DOM will be ambiguous.
-  if (frame_url.is_empty()) {
-    ambiguous_dom_ = true;
-  }
-
   // Add the urls from the DOM to |resources_|. The renderer could be sending
   // bogus messages, so limit the number of nodes we accept.
   // Also update |elements_| with the DOM structure.
@@ -510,13 +473,23 @@ void ThreatDetails::AddDOMDetails(
     }
     // Check for a tag_name to avoid adding the summary node to the DOM.
     if (!node.tag_name.empty()) {
-      AddDomElement(frame_tree_node_id, frame_url.spec(), node.node_id,
-                    node.tag_name, node.parent_node_id, node.attributes,
-                    resource);
+      AddDomElement(process_id, frame_tree_node_id, node.other_frame_routing_id,
+                    node.node_id, node.tag_name, node.parent_node_id,
+                    node.attributes, resource);
     }
   }
 }
 
+void ThreatDetails::LookupOtherFrameId(const std::string& element_key,
+                                       const int process_id,
+                                       const int other_frame_routing_id) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+  int other_frame_tree_node_id =
+      content::RenderFrameHost::LookupOtherFrameTreeNodeId(
+          process_id, other_frame_routing_id);
+  iframe_key_to_frame_tree_id_map_[element_key] = other_frame_tree_node_id;
+}
+
 // Called from the SB Service on the IO thread, after the user has
 // closed the tab, or clicked proceed or goback.  Since the user needs
 // to take an action, we expect this to be called after
@@ -525,6 +498,22 @@ void ThreatDetails::AddDOMDetails(
 void ThreatDetails::FinishCollection(bool did_proceed, int num_visit) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
 
+  // Do a second pass over the elements and update iframe elements to have
+  // references to their children. Children will have been received from a
+  // different renderer than the iframe element.
+  for (auto& element_pair : elements_) {
+    const std::string& element_key = element_pair.first;
+    HTMLElement* element = element_pair.second.get();
+    if (element->tag() == "IFRAME" || element->tag() == "FRAME") {
+      int frame_tree_id_of_iframe_renderer =
+          iframe_key_to_frame_tree_id_map_[element_key];
+      const std::unordered_set<int>& child_ids =
+          frame_tree_id_to_children_map_[frame_tree_id_of_iframe_renderer];
+      for (const int child_id : child_ids) {
+        element->add_child_ids(child_id);
+      }
+    }
+  }
   did_proceed_ = did_proceed;
   num_visits_ = num_visit;
   std::vector<GURL> urls;
@@ -576,12 +565,6 @@ void ThreatDetails::OnCacheCollectionReady() {
   for (auto& element_pair : elements_) {
     report_->add_dom()->Swap(element_pair.second.get());
   }
-  if (!elements_.empty()) {
-    // TODO(lpz): Consider including the ambiguous_dom_ bit in the report
-    // itself.
-    UMA_HISTOGRAM_BOOLEAN("SafeBrowsing.ThreatReport.DomIsAmbiguous",
-                          ambiguous_dom_);
-  }
 
   report_->set_did_proceed(did_proceed_);
   // Only sets repeat_visit if num_visits_ >= 0.