Remove Dangerous indicator after going back from interstitial

chrome/browser/safe_browsing/ui_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 "chrome/browser/safe_browsing/ui_manager.h"
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/callback.h"
 #include "base/debug/leak_tracker.h"
@@ skipping 28 matching lines @@
 using content::NavigationEntry;
 using content::WebContents;
 using safe_browsing::HitReport;
 
 namespace {
 
 const void* const kWhitelistKey = &kWhitelistKey;
 
 // A WhitelistUrlSet holds the set of URLs that have been whitelisted for a
 // specific WebContents, along with pending entries that are still undecided.
+// The URLs in this set should come from GetWhitelistUrl() or
+// GetMainFrameWhitelistUrlForResource().
 class WhitelistUrlSet : public base::SupportsUserData::Data {
  public:
   WhitelistUrlSet() {}
 
-  bool Contains(const GURL url) {
-    return set_.find(url.GetWithEmptyPath()) != set_.end();
-  }
+  bool Contains(const GURL url) { return set_.find(url) != set_.end(); }
+
+  void RemovePending(const GURL& url) { pending_.erase(url); }
 
   void Insert(const GURL url) {
-    set_.insert(url.GetWithEmptyPath());
-    pending_.erase(url.GetWithEmptyPath());
+    set_.insert(url);
+    RemovePending(url);
   }
 
   bool ContainsPending(const GURL url) {
-    return pending_.find(url.GetWithEmptyPath()) != pending_.end();
+    return pending_.find(url) != pending_.end();
   }
 
-  void InsertPending(const GURL url) {
-    pending_.insert(url.GetWithEmptyPath());
-  }
+  void InsertPending(const GURL url) { pending_.insert(url); }
 
  private:
   std::set<GURL> set_;
   std::set<GURL> pending_;
 
   DISALLOW_COPY_AND_ASSIGN(WhitelistUrlSet);
 };
 
+// Returns the URL that should be used in a WhitelistUrlSet for the given
+// |resource|.
+GURL GetMainFrameWhitelistUrlForResource(
+    const safe_browsing::SafeBrowsingUIManager::UnsafeResource& resource) {
+  if (resource.is_subresource) {
+    NavigationEntry* entry = resource.GetNavigationEntryForResource();
+    if (!entry)
+      return GURL();
+    return entry->GetURL().GetWithEmptyPath();
+  }
+  return resource.url.GetWithEmptyPath();
+}
+
+// Returns the URL that should be used in a WhitelistUrlSet for the
+// resource loaded from |url| on a navigation |entry|.
+GURL GetWhitelistUrl(const GURL& url,
+                     bool is_subresource,
+                     NavigationEntry* entry) {
+  if (is_subresource) {
+    if (!entry)
+      return GURL();
+    return entry->GetURL().GetWithEmptyPath();
+  }
+  return url.GetWithEmptyPath();
+}
+
+WhitelistUrlSet* GetOrCreateWhitelist(content::WebContents* web_contents) {
+  WhitelistUrlSet* site_list =
+      static_cast<WhitelistUrlSet*>(web_contents->GetUserData(kWhitelistKey));
+  if (!site_list) {
+    site_list = new WhitelistUrlSet;
+    web_contents->SetUserData(kWhitelistKey, site_list);
+  }
+  return site_list;
+}
+
 }  // namespace
 
 namespace safe_browsing {
 
 // SafeBrowsingUIManager::UnsafeResource ---------------------------------------
 
 SafeBrowsingUIManager::UnsafeResource::UnsafeResource()
     : is_subresource(false),
       threat_type(SB_THREAT_TYPE_SAFE),
       threat_source(safe_browsing::ThreatSource::UNKNOWN) {}
@@ skipping 58 matching lines @@
   if (shutdown)
     sb_service_ = NULL;
 }
 
 void SafeBrowsingUIManager::LogPauseDelay(base::TimeDelta time) {
   UMA_HISTOGRAM_LONG_TIMES("SB2.Delay", time);
 }
 
 void SafeBrowsingUIManager::OnBlockingPageDone(
     const std::vector<UnsafeResource>& resources,
-    bool proceed) {
+    bool proceed,
+    content::WebContents* web_contents,
+    const GURL& main_frame_url) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   for (const auto& resource : resources) {
     if (!resource.callback.is_null()) {
       DCHECK(resource.callback_thread);
       resource.callback_thread->PostTask(
           FROM_HERE, base::Bind(resource.callback, proceed));
     }
 
-    if (proceed)
-      AddToWhitelistUrlSet(resource, false /* Pending -> permanent */);
+    GURL whitelist_url = GetWhitelistUrl(
+        main_frame_url, false /* is subresource */,
+        nullptr /* no navigation entry needed for main resource */);
+    if (proceed) {
+      AddToWhitelistUrlSet(whitelist_url, web_contents,
+                           false /* Pending -> permanent */);
+    } else if (web_contents) {
+      // |web_contents| doesn't exist if the tab has been closed.
+      RemoveFromPendingWhitelistUrlSet(whitelist_url, web_contents);
+    }
   }
 }
 
 void SafeBrowsingUIManager::DisplayBlockingPage(
     const UnsafeResource& resource) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   if (resource.is_subresource && !resource.is_subframe) {
     // Sites tagged as serving Unwanted Software should only show a warning for
     // main-frame or sub-frame resource. Similar warning restrictions should be
     // applied to malware sites tagged as "landing sites" (see "Types of
@@ skipping 13 matching lines @@
       return;
     }
   }
 
   // The tab might have been closed. If it was closed, just act as if "Don't
   // Proceed" had been chosen.
   WebContents* web_contents = resource.web_contents_getter.Run();
   if (!web_contents) {
     std::vector<UnsafeResource> resources;
     resources.push_back(resource);
-    OnBlockingPageDone(resources, false);
+    OnBlockingPageDone(resources, false, web_contents,
+                       GetMainFrameWhitelistUrlForResource(resource));
     return;
   }
 
   // Check if the user has already ignored a SB warning for the same WebContents
   // and top-level domain.
   if (IsWhitelisted(resource)) {
     if (!resource.callback.is_null()) {
       DCHECK(resource.callback_thread);
       resource.callback_thread->PostTask(FROM_HERE,
                                          base::Bind(resource.callback, true));
@@ skipping 35 matching lines @@
     hit_report.is_metrics_reporting_active =
         ChromeMetricsServiceAccessor::IsMetricsAndCrashReportingEnabled();
 
     MaybeReportSafeBrowsingHit(hit_report);
   }
 
   if (resource.threat_type != SB_THREAT_TYPE_SAFE) {
     for (Observer& observer : observer_list_)
       observer.OnSafeBrowsingHit(resource);
   }
-  AddToWhitelistUrlSet(resource, true /* A decision is now pending */);
+  AddToWhitelistUrlSet(GetMainFrameWhitelistUrlForResource(resource),
+                       resource.web_contents_getter.Run(),
+                       true /* A decision is now pending */);
   SafeBrowsingBlockingPage::ShowBlockingPage(this, resource);
 }
 
 // A safebrowsing hit is sent after a blocking page for malware/phishing
 // or after the warning dialog for download urls, only for
 // UMA || extended_reporting users.
 void SafeBrowsingUIManager::MaybeReportSafeBrowsingHit(
     const HitReport& hit_report) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
 
@@ skipping 45 matching lines @@
 void SafeBrowsingUIManager::AddObserver(Observer* observer) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   observer_list_.AddObserver(observer);
 }
 
 void SafeBrowsingUIManager::RemoveObserver(Observer* observer) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   observer_list_.RemoveObserver(observer);
 }
 
+// Static.
+void SafeBrowsingUIManager::CreateWhitelistForTesting(
+    content::WebContents* web_contents) {
+  GetOrCreateWhitelist(web_contents);
+}
+
 void SafeBrowsingUIManager::ReportInvalidCertificateChainOnIOThread(
     const std::string& serialized_report) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
 
   // The service may delete the ping manager (i.e. when user disabling service,
   // etc). This happens on the IO thread.
   if (!sb_service_ || !sb_service_->ping_manager())
     return;
 
   sb_service_->ping_manager()->ReportInvalidCertificateChain(serialized_report);
@@ skipping 21 matching lines @@
   // etc). This happens on the IO thread.
   if (sb_service_.get() == NULL || sb_service_->ping_manager() == NULL)
     return;
 
   if (!serialized.empty()) {
     DVLOG(1) << "Sending serialized threat details.";
     sb_service_->ping_manager()->ReportThreatDetails(serialized);
   }
 }
 
-// Record this domain in the current WebContents as either whitelisted or
+// Record this domain in the given WebContents as either whitelisted or
 // pending whitelisting (if an interstitial is currently displayed). If an
 // existing WhitelistUrlSet does not yet exist, create a new WhitelistUrlSet.
-void SafeBrowsingUIManager::AddToWhitelistUrlSet(const UnsafeResource& resource,
-                                                 bool pending) {
+void SafeBrowsingUIManager::AddToWhitelistUrlSet(
+    const GURL& whitelist_url,
+    content::WebContents* web_contents,
+    bool pending) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
 
-  WebContents* web_contents = resource.web_contents_getter.Run();
+  // A WebContents might not exist if the tab has been closed.
+  if (!web_contents)
+    return;
+
+  WhitelistUrlSet* site_list = GetOrCreateWhitelist(web_contents);
+
+  if (whitelist_url.is_empty())
+    return;
+
+  if (pending) {
+    site_list->InsertPending(whitelist_url);
+  } else {
+    site_list->Insert(whitelist_url);
+  }
+
+  // Notify security UI that security state has changed.
+  web_contents->DidChangeVisibleSecurityState();
+}
+
+void SafeBrowsingUIManager::RemoveFromPendingWhitelistUrlSet(
+    const GURL& whitelist_url,
+    content::WebContents* web_contents) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+
+  // A WebContents might not exist if the tab has been closed.
+  if (!web_contents)
+    return;
+
+  // Use |web_contents| rather than |resource.web_contents_getter|
+  // here. By this point, a "Back" navigation could have already been
+  // committed, so the page loading |resource| might be gone and
+  // |web_contents_getter| may no longer be valid.
   WhitelistUrlSet* site_list =
       static_cast<WhitelistUrlSet*>(web_contents->GetUserData(kWhitelistKey));
-  if (!site_list) {
-    site_list = new WhitelistUrlSet;
-    web_contents->SetUserData(kWhitelistKey, site_list);
-  }
 
-  GURL whitelisted_url;
-  if (resource.is_subresource) {
-    NavigationEntry* entry = resource.GetNavigationEntryForResource();
-    if (!entry)
-      return;
-    whitelisted_url = entry->GetURL();
-  } else {
-    whitelisted_url = resource.url;
-  }
+  if (whitelist_url.is_empty())
+    return;
 
-  if (pending) {
-    site_list->InsertPending(whitelisted_url);
-  } else {
-    site_list->Insert(whitelisted_url);
-  }
+  // Note that this function does not DCHECK that |whitelist_url|
+  // appears in the pending whitelist. In the common case, it's expected
+  // that a URL is in the pending whitelist when it is removed, but it's
+  // not always the case. For example, if there are several blocking
+  // pages queued up for different resources on the same page, and the
+  // user goes back to dimiss the first one, the subsequent blocking
+  // pages get dismissed as well (as if the user had clicked "Back to
+  // safety" on each of them). In this case, the first dismissal will
+  // remove the main-frame URL from the pending whitelist, so the
+  // main-frame URL will have already been removed when the subsequent
+  // blocking pages are dismissed.
+  if (site_list->ContainsPending(whitelist_url))
+    site_list->RemovePending(whitelist_url);
 
   // Notify security UI that security state has changed.
   web_contents->DidChangeVisibleSecurityState();
 }
 
 bool SafeBrowsingUIManager::IsWhitelisted(const UnsafeResource& resource) {
   NavigationEntry* entry = nullptr;
   if (resource.is_subresource) {
     entry = resource.GetNavigationEntryForResource();
   }
   return IsUrlWhitelistedOrPendingForWebContents(
       resource.url, resource.is_subresource, entry,
       resource.web_contents_getter.Run(), true);
 }
 
 // Check if the user has already seen and/or ignored a SB warning for this
 // WebContents and top-level domain.
 bool SafeBrowsingUIManager::IsUrlWhitelistedOrPendingForWebContents(
     const GURL& url,
     bool is_subresource,
     NavigationEntry* entry,
     content::WebContents* web_contents,
     bool whitelist_only) {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
 
-  GURL lookup_url;
-  if (is_subresource) {
-    if (!entry)
-      return false;
-    lookup_url = entry->GetURL();
-  } else {
-    lookup_url = url;
-  }
+  GURL lookup_url = GetWhitelistUrl(url, is_subresource, entry);
+  if (lookup_url.is_empty())
+    return false;
 
   WhitelistUrlSet* site_list =
       static_cast<WhitelistUrlSet*>(web_contents->GetUserData(kWhitelistKey));
   if (!site_list)
     return false;
 
   bool whitelisted = site_list->Contains(lookup_url);
   if (whitelist_only) {
     return whitelisted;
   } else {
     return whitelisted || site_list->ContainsPending(lookup_url);
   }
 }
 
+// Static.
+GURL SafeBrowsingUIManager::GetMainFrameWhitelistUrlForResourceForTesting(
+    const safe_browsing::SafeBrowsingUIManager::UnsafeResource& resource) {
+  return GetMainFrameWhitelistUrlForResource(resource);
+}
+
 }  // namespace safe_browsing