Remove Dangerous indicator after going back from interstitial

chrome/browser/safe_browsing/ui_manager.cc

Index: chrome/browser/safe_browsing/ui_manager.cc
diff --git a/chrome/browser/safe_browsing/ui_manager.cc b/chrome/browser/safe_browsing/ui_manager.cc
index 6cdc91586cff1ea2f834bfce8828ab1337c72777..bcc9f655f9bbeb5dd50932c6c30541a8f7bb0cd7 100644
--- a/chrome/browser/safe_browsing/ui_manager.cc
+++ b/chrome/browser/safe_browsing/ui_manager.cc
@@ -46,26 +46,26 @@ 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_;
@@ -74,6 +74,42 @@ class WhitelistUrlSet : public base::SupportsUserData::Data {
   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 {
@@ -152,7 +188,9 @@ void SafeBrowsingUIManager::LogPauseDelay(base::TimeDelta 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()) {
@@ -161,8 +199,16 @@ void SafeBrowsingUIManager::OnBlockingPageDone(
           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);
+    }
   }
 }
 
@@ -196,7 +242,8 @@ void SafeBrowsingUIManager::DisplayBlockingPage(
   if (!web_contents) {
     std::vector<UnsafeResource> resources;
     resources.push_back(resource);
-    OnBlockingPageDone(resources, false);
+    OnBlockingPageDone(resources, false, web_contents,
+                       GetMainFrameWhitelistUrlForResource(resource));
     return;
   }
 
@@ -252,7 +299,9 @@ void SafeBrowsingUIManager::DisplayBlockingPage(
     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);
 }
 
@@ -318,6 +367,12 @@ void SafeBrowsingUIManager::RemoveObserver(Observer* observer) {
   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);
@@ -359,41 +414,71 @@ void SafeBrowsingUIManager::SendSerializedThreatDetails(
   }
 }
 
-// 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();
-  WhitelistUrlSet* site_list =
-      static_cast<WhitelistUrlSet*>(web_contents->GetUserData(kWhitelistKey));
-  if (!site_list) {
-    site_list = new WhitelistUrlSet;
-    web_contents->SetUserData(kWhitelistKey, site_list);
-  }
+  // A WebContents might not exist if the tab has been closed.
+  if (!web_contents)
+    return;
 
-  GURL whitelisted_url;
-  if (resource.is_subresource) {
-    NavigationEntry* entry = resource.GetNavigationEntryForResource();
-    if (!entry)
-      return;
-    whitelisted_url = entry->GetURL();
-  } else {
-    whitelisted_url = resource.url;
-  }
+  WhitelistUrlSet* site_list = GetOrCreateWhitelist(web_contents);
+
+  if (whitelist_url.is_empty())
+    return;
 
   if (pending) {
-    site_list->InsertPending(whitelisted_url);
+    site_list->InsertPending(whitelist_url);
   } else {
-    site_list->Insert(whitelisted_url);
+    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 (whitelist_url.is_empty())
+    return;
+
+  // 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) {
@@ -414,14 +499,9 @@ bool SafeBrowsingUIManager::IsUrlWhitelistedOrPendingForWebContents(
     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));
@@ -436,4 +516,10 @@ bool SafeBrowsingUIManager::IsUrlWhitelistedOrPendingForWebContents(
   }
 }
 
+// Static.
+GURL SafeBrowsingUIManager::GetMainFrameWhitelistUrlForResourceForTesting(
+    const safe_browsing::SafeBrowsingUIManager::UnsafeResource& resource) {
+  return GetMainFrameWhitelistUrlForResource(resource);
+}
+
 }  // namespace safe_browsing