Add a Node message in the malware details protocol buffer. This allows us to ...

chrome/browser/safe_browsing/malware_details.cc

Index: chrome/browser/safe_browsing/malware_details.cc
===================================================================
--- chrome/browser/safe_browsing/malware_details.cc	(revision 71096)
+++ chrome/browser/safe_browsing/malware_details.cc	(working copy)
@@ -11,6 +11,8 @@
 #include "chrome/browser/tab_contents/tab_contents.h"
 #include "chrome/browser/safe_browsing/report.pb.h"
 
+using safe_browsing::ClientMalwareReportRequest;
+
 // Create a MalwareDetails for the given tab. Runs in the UI thread.
 MalwareDetails::MalwareDetails(
     TabContents* tab_contents,
@@ -26,21 +28,66 @@
   return url.SchemeIs("http");  // TODO(panayiotis): also skip internal urls.
 }
 
-void MalwareDetails::AddNode(const std::string& url,
+// Looks for a Resource and Node protocol messages for the given url
+// in resources_ and nodes_.  If found, it updates |resource| and
+// |node|. Otherwise, it creates new messages, adds them to resources_
+// and nodes_, and updates |resource| and |node| to point to them.
+// Note: We create exactly one Node for each Resource.
+void MalwareDetails::FindOrCreateResource(
+    const std::string& url,
+    ClientMalwareReportRequest::Resource** resource,
+    ClientMalwareReportRequest::Node** node) {
+  ResourceMap::iterator it = resources_.find(url);
+  if (it != resources_.end()) {
+    // We have the resource, also find the corresponding node.
+    *resource = it->second.get();
+    NodeMap::iterator node_it = nodes_.find((*resource)->id());
+    CHECK(node_it != nodes_.end());
+    *node = node_it->second.get();
+  } else {
+    // Create the resource and node for |url|
+    int id = resources_.size();
+    linked_ptr<ClientMalwareReportRequest::Resource> new_resource(
+        new ClientMalwareReportRequest::Resource());
+    new_resource->set_url(url);
+    new_resource->set_id(id);
+    resources_[url] = new_resource;
+    *resource = new_resource.get();
+
+    linked_ptr<ClientMalwareReportRequest::Node> new_node(
+        new ClientMalwareReportRequest::Node());
+    new_node->set_id(id);
+    nodes_[id] = new_node;
+    *node = new_node.get();
+  }
+}
+
+void MalwareDetails::AddUrl(const std::string& url,
                              const std::string& parent) {
   if (!IsPublicUrl(GURL(url)))
     return;
-  linked_ptr<safe_browsing::ClientMalwareReportRequest::Resource> resource(
-      new safe_browsing::ClientMalwareReportRequest::Resource());
-  resource->set_url(url);
-  if (!parent.empty() && IsPublicUrl(GURL(parent)))
-    resource->set_parent(parent);
-  urls_[url] = resource;
+
+  // Find (or create) the resource for the url.
+  ClientMalwareReportRequest::Resource* url_resource = NULL;
+  ClientMalwareReportRequest::Node* url_node = NULL;
+  FindOrCreateResource(url, &url_resource, &url_node);

[lzheng, 2011/01/12 22:30:15]

Can we break this function to to? Find() and Create()?

[panayiotis, 2011/01/13 02:25:00]

It's now quite shorter, do you think it still needs to be broken into two? 

+  CHECK(url_resource != NULL && url_node != NULL);
+
+  if (!parent.empty() && IsPublicUrl(GURL(parent))) {
+    // Add the resource and node for the parent node.
+    ClientMalwareReportRequest::Resource* parent_resource = NULL;
+    ClientMalwareReportRequest::Node* parent_node = NULL;
+    FindOrCreateResource(parent, &parent_resource, &parent_node);
+    CHECK(parent_resource != NULL && parent_node != NULL);
+
+    // Update the parent-child relation
+    url_node->set_parent_id(parent_node->id());
+  }
 }
 
 void MalwareDetails::StartCollection() {
   DVLOG(1) << "Starting to compute malware details.";
-  report_.reset(new safe_browsing::ClientMalwareReportRequest());
+  report_.reset(new ClientMalwareReportRequest());
 
   if (IsPublicUrl(resource_.url)) {
     report_->set_malware_url(resource_.url.spec());
@@ -61,30 +108,35 @@
   }
 
   // Add the nodes, starting from the page url.
-  AddNode(page_url.spec(), "");
+  AddUrl(page_url.spec(), "");
 
   // Add the resource_url and its original url, if non-empty and different.
   if (!resource_.original_url.spec().empty() &&
       resource_.url != resource_.original_url) {
     // Add original_url, as the parent of resource_url.
-    AddNode(resource_.original_url.spec(), "");
-    AddNode(resource_.url.spec(), resource_.original_url.spec());
+    AddUrl(resource_.original_url.spec(), "");
+    AddUrl(resource_.url.spec(), resource_.original_url.spec());
   } else {
-    AddNode(resource_.url.spec(), "");
+    AddUrl(resource_.url.spec(), "");
   }
 
   // Add the referrer url.
   if (nav_entry && !referrer_url.spec().empty()) {
-    AddNode(referrer_url.spec(), "");
+    AddUrl(referrer_url.spec(), "");
   }
 
-  // Add all the urls in our |urls_| map to the |report_| protobuf.
-  for (ResourceMap::const_iterator it = urls_.begin();
-       it != urls_.end(); it++) {
-    safe_browsing::ClientMalwareReportRequest::Resource* pb_resource =
-        report_->add_nodes();
+  // The |report_| protocol buffer is now generated: We add all the
+  // urls in our |resources_| and |nodes_| maps.
+  for (ResourceMap::const_iterator it = resources_.begin();
+       it != resources_.end(); it++) {
+    ClientMalwareReportRequest::Resource* pb_resource =
+        report_->add_resources();
     pb_resource->CopyFrom(*(it->second));
   }
+  for (NodeMap::const_iterator it = nodes_.begin(); it != nodes_.end(); it++) {
+    ClientMalwareReportRequest::Node* pb_node = report_->add_nodes();
+    pb_node->CopyFrom(*(it->second));
+  }
 }
 
 // Called from the SB Service on the IO thread.