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

chrome/browser/safe_browsing/malware_details.cc

 // Copyright (c) 2010 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.
 //
 // Implementation of the MalwareDetails class.
 
 #include "chrome/browser/safe_browsing/malware_details.h"
 
 #include "chrome/browser/safe_browsing/safe_browsing_service.h"
 #include "chrome/browser/tab_contents/navigation_entry.h"
 #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,
     const SafeBrowsingService::UnsafeResource resource)
     : tab_contents_(tab_contents),
       resource_(resource) {
   StartCollection();
 }
 
 MalwareDetails::~MalwareDetails() {}
 
 bool MalwareDetails::IsPublicUrl(const GURL& url) const {
   return url.SchemeIs("http");  // TODO(panayiotis): also skip internal urls.
 }
 
-void MalwareDetails::AddNode(const std::string& url,
+// Looks for a Resource for the given url in resources_.  If found, it
+// updates |resource|. Otherwise, it creates a new message, adds it to
+// resources_ and updates |resource| to point to it.
+ClientMalwareReportRequest::Resource* MalwareDetails::FindOrCreateResource(
+    const std::string& url) {
+  ResourceMap::iterator it = resources_.find(url);
+  if (it != resources_.end()) {
+    return it->second.get();
+  }
+
+  // Create the resource 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;
+  return new_resource.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 =
+      FindOrCreateResource(url);
+  if (!parent.empty() && IsPublicUrl(GURL(parent))) {
+    // Add the resource for the parent.
+    ClientMalwareReportRequest::Resource* parent_resource =
+        FindOrCreateResource(parent);
+    // Update the parent-child relation
+    url_resource->set_parent_id(parent_resource->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());
   }
 
   GURL page_url = tab_contents_->GetURL();
   if (IsPublicUrl(page_url)) {
     report_->set_page_url(page_url.spec());
   }
 
   GURL referrer_url;
   NavigationEntry* nav_entry = tab_contents_->controller().GetActiveEntry();
   if (nav_entry) {
     referrer_url = nav_entry->referrer();
     if (IsPublicUrl(referrer_url)) {
       report_->set_referrer_url(referrer_url.spec());
     }
   }
 
   // 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_| maps.
+  for (ResourceMap::const_iterator it = resources_.begin();
+       it != resources_.end(); it++) {
+    ClientMalwareReportRequest::Resource* pb_resource =
+        report_->add_resources();
     pb_resource->CopyFrom(*(it->second));
   }
 }
 
 // Called from the SB Service on the IO thread.
 const std::string* MalwareDetails::GetSerializedReport() {
   scoped_ptr<std::string> request_data(new std::string());
   if (!report_->SerializeToString(request_data.get())) {
     DLOG(ERROR) << "Unable to serialize the malware report.";
   }
 
   return request_data.release();
 }