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 27 matching lines @@
 using content::BrowserThread;
 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.

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

Add a comment that the URLs here should come from GetWhitelistUrl()

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

Done.

 class WhitelistUrlSet : public base::SupportsUserData::Data {
  public:
   WhitelistUrlSet() {}
 
   bool Contains(const GURL url) {
     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) {
     return pending_.find(url.GetWithEmptyPath()) != pending_.end();
   }
 
   void InsertPending(const GURL url) {
     pending_.insert(url.GetWithEmptyPath());
   }
 
  private:
   std::set<GURL> set_;
   std::set<GURL> pending_;
 
   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 {
 
 // SafeBrowsingUIManager::UnsafeResource ---------------------------------------
 
 SafeBrowsingUIManager::UnsafeResource::UnsafeResource()
     : is_subresource(false),
       threat_type(SB_THREAT_TYPE_SAFE),
       threat_source(safe_browsing::ThreatSource::UNKNOWN) {}
@@ skipping 69 matching lines @@
   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 */);
+    else
+      RemoveFromPendingWhitelistUrlSet(resource);
   }
 }
 
 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 191 matching lines @@
   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);
   }
 
-  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);
   } else {
     site_list->Insert(whitelisted_url);
   }
 
   // Notify security UI that security state has changed.
   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) {
     entry = resource.GetNavigationEntryForResource();
   }
   return IsUrlWhitelistedOrPendingForWebContents(
       resource.url, resource.is_subresource, entry,
       resource.web_contents_getter.Run(), true);
 }
 
@@ skipping 23 matching lines @@
 
   bool whitelisted = site_list->Contains(lookup_url);
   if (whitelist_only) {
     return whitelisted;
   } else {
     return whitelisted || site_list->ContainsPending(lookup_url);
   }
 }
 
 }  // namespace safe_browsing