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..aa96596d648e1e63ca942ebcc39264ec52f76174 100644
--- a/chrome/browser/safe_browsing/ui_manager.cc
+++ b/chrome/browser/safe_browsing/ui_manager.cc
@@ -54,9 +54,11 @@ class WhitelistUrlSet : public base::SupportsUserData::Data {
     return set_.find(url.GetWithEmptyPath()) != set_.end();
   }
 
+  void RemovePending(const GURL url) { pending_.erase(url.GetWithEmptyPath()); }

[Nathan Parker, 2016/10/27 18:31:02]

const GURL&
Though maybe the compiler would do that for you since it's inline -- not sure.

[estark, 2016/10/28 17:57:47]

Done.

+
   void Insert(const GURL url) {
     set_.insert(url.GetWithEmptyPath());
-    pending_.erase(url.GetWithEmptyPath());
+    RemovePending(url);
   }
 
   bool ContainsPending(const GURL url) {
@@ -74,6 +76,17 @@ class WhitelistUrlSet : public base::SupportsUserData::Data {
   DISALLOW_COPY_AND_ASSIGN(WhitelistUrlSet);
 };
 
+GURL GetWhitelistUrl(
+    const safe_browsing::SafeBrowsingUIManager::UnsafeResource& resource) {
+  if (resource.is_subresource) {
+    NavigationEntry* entry = resource.GetNavigationEntryForResource();
+    if (!entry)
+      return GURL();
+    return entry->GetURL();

[Nathan Parker, 2016/10/27 18:31:02]

How about returning .GetWithEmptyPath() so you don't have to do it repeatedly
above?  That'll also ensure you don't forget it some place above.

[estark, 2016/10/28 17:57:47]

Done.

+  }
+  return resource.url;
+}
+
 }  // namespace
 
 namespace safe_browsing {
@@ -163,6 +176,8 @@ void SafeBrowsingUIManager::OnBlockingPageDone(
 
     if (proceed)
       AddToWhitelistUrlSet(resource, false /* Pending -> permanent */);
+    else
+      RemoveFromPendingWhitelistUrlSet(resource);
   }
 }
 
@@ -374,15 +389,9 @@ void SafeBrowsingUIManager::AddToWhitelistUrlSet(const UnsafeResource& resource,
     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;
-  }
+  GURL whitelisted_url = GetWhitelistUrl(resource);
+  if (whitelisted_url.is_empty())
+    return;
 
   if (pending) {
     site_list->InsertPending(whitelisted_url);
@@ -394,6 +403,21 @@ void SafeBrowsingUIManager::AddToWhitelistUrlSet(const UnsafeResource& resource,
   web_contents->DidChangeVisibleSecurityState();
 }
 
+void SafeBrowsingUIManager::RemoveFromPendingWhitelistUrlSet(
+    const UnsafeResource& resource) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+  WebContents* web_contents = resource.web_contents_getter.Run();
+  WhitelistUrlSet* site_list =
+      static_cast<WhitelistUrlSet*>(web_contents->GetUserData(kWhitelistKey));
+  GURL whitelisted_url = GetWhitelistUrl(resource);
+  if (whitelisted_url.is_empty())
+    return;
+  DCHECK(site_list->ContainsPending(whitelisted_url));
+  site_list->RemovePending(whitelisted_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) {