Content API changes to improve DOM stitching in ThreatDetails code.

components/safe_browsing/browser/threat_details.h

Index: components/safe_browsing/browser/threat_details.h
diff --git a/components/safe_browsing/browser/threat_details.h b/components/safe_browsing/browser/threat_details.h
index 882458b64d5ad75aa9b6fa4c0e72dedb00100ad0..16a244f9e222ca95e4c33abb338b5ca1a1916b19 100644
--- a/components/safe_browsing/browser/threat_details.h
+++ b/components/safe_browsing/browser/threat_details.h
@@ -52,22 +52,17 @@ using ResourceMap =
 
 // Maps a key of an HTML element to its corresponding HTMLElement proto message.
 // HTML Element keys have the form "<frame_id>-<node_id>", where |frame_id| is
-// the FrameTree NodeID of the render frame containing the element, and
+// the FrameTreeNode ID of the frame containing the element, and
 // |node_id| is a sequential ID for the element generated by the renderer.
 using ElementMap = base::hash_map<std::string, std::unique_ptr<HTMLElement>>;
 
-// Maps a URL to some HTML Elements. Used to maintain parent/child relationship
-// for HTML Elements across IFrame boundaries.
-// The key is the string URL set as the src attribute of an iframe. The value is
-// the HTMLElement proto that represents the iframe element with that URL.
-// The HTMLElement protos are not owned by this map.
-using UrlToDomElementMap = base::hash_map<std::string, HTMLElement*>;
+// Maps the key of an iframe element to the FrameTreeNode ID of the frame that
+// rendered the contents of the iframe.
+using KeyToFrameTreeIdMap = base::hash_map<std::string, int>;
 
-// Maps a URL to some Element IDs. Used to maintain parent/child relationship
-// for HTML Elements across IFrame boundaries.
-// The key is the string URL of a render frame. The value is the set of Element
-// IDs that are at the top-level of this render frame.
-using UrlToChildIdsMap = base::hash_map<std::string, std::unordered_set<int>>;
+// Maps a FrameTreeNode ID of a frame to a set of child IDs. The child IDs are
+// the Element IDs of the top-level HTML Elements in this frame.
+using FrameTreeIdToChildIdsMap = base::hash_map<int, std::unordered_set<int>>;
 
 class ThreatDetails : public base::RefCountedThreadSafe<
                           ThreatDetails,
@@ -121,8 +116,8 @@ class ThreatDetails : public base::RefCountedThreadSafe<
   // Called on the IO thread with the DOM details.
   virtual void AddDOMDetails(
       const int frame_tree_node_id,
-      const GURL& frame_last_committed_url,
-      const std::vector<SafeBrowsingHostMsg_ThreatDOMDetails_Node>& params);
+      const std::vector<SafeBrowsingHostMsg_ThreatDOMDetails_Node>& params,
+      const KeyToFrameTreeIdMap& child_frame_tree_map);
 
   // The report protocol buffer.
   std::unique_ptr<ClientSafeBrowsingReportRequest> report_;
@@ -167,16 +162,13 @@ class ThreatDetails : public base::RefCountedThreadSafe<
 
   void AddRedirectUrlList(const std::vector<GURL>& urls);
 
-  // Adds an HTML Element to the DOM structure.
-  // |frame_tree_node_id| is the unique ID of the render frame the element came
-  // from. |frame_url| is the URL that the render frame was handling.
-  // |element_node_id| is a unique ID of the element within the render frame.
-  // |tag_name| is the tag of the element. |parent_element_node_id| is the
-  // unique ID of the parent element with the render frame. |attributes|
-  // contains the names and values of the element's attributes.|resource| is
-  // set if this element is a resource.
+  // Adds an HTML Element to the DOM structure. |frame_tree_node_id| is the
+  // unique ID of the frame the element came from. |element_node_id| is a unique
+  // ID of the element within the frame. |tag_name| is the tag of the element.
+  // |parent_element_node_id| is the unique ID of the parent element within the
+  // frame. |attributes| contains the names and values of the element's
+  // attributes. |resource| is set if this element is a resource.
   void AddDomElement(const int frame_tree_node_id,
-                     const std::string& frame_url,
                      const int element_node_id,
                      const std::string& tag_name,
                      const int parent_element_node_id,
@@ -194,18 +186,20 @@ class ThreatDetails : public base::RefCountedThreadSafe<
   // Store all HTML elements collected, keep them in a map for easy lookup.
   ElementMap elements_;
 
-  // For each iframe element encountered we map the src of the iframe to the
-  // iframe element. This is used when we receive elements from a different
-  // frame whose document URL matches the src of an iframe in this map. We can
-  // then add all elements from the subframe as children of the iframe element
-  // stored here.
-  UrlToDomElementMap iframe_src_to_element_map_;
-
-  // When getting a set of elements from a render frame, we store the frame's
-  // URL and a collection of all the top-level elements in that frame. When we
-  // later encounter the parent iframe with the same src URL, we can add all of
-  // these elements as children of that iframe.
-  UrlToChildIdsMap document_url_to_children_map_;
+  // For each iframe element encountered we map the key of the iframe to the
+  // FrameTreeNode ID of the frame containing the contents of that iframe.
+  // We populate this map when receiving results from ThreatDomDetails, and use
+  // it in a second pass (after FinishCollection) to attach children to iframe
+  // elements.
+  // Should only be accessed on the IO thread.
+  KeyToFrameTreeIdMap iframe_key_to_frame_tree_id_map_;
+
+  // When getting a set of elements from a frame, we store the frame's
+  // FrameTreeNode ID and a collection of all top-level elements in that frame.
+  // It is populated as we receive sets of nodes from different renderers.
+  // It is used together with |iframe_key_to_frame_tree_id_map_| in a second
+  // pass to insert child elements under their parent iframe elements.
+  FrameTreeIdToChildIdsMap frame_tree_id_to_children_map_;
 
   // Result from the cache extractor.
   bool cache_result_;
@@ -218,7 +212,7 @@ class ThreatDetails : public base::RefCountedThreadSafe<
   int num_visits_;
 
   // Keeps track of whether we have an ambiguous DOM in this report. This can
-  // happen when the HTML Elements returned by a render frame can't be
+  // happen when the HTML Elements returned by a renderer can't be
   // associated with a parent Element in the parent frame.
   bool ambiguous_dom_;