Fetch extension blacklist states from SafeBrowsing server

chrome/browser/extensions/blacklist_unittest.cc

 // Copyright 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 "base/bind.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "chrome/browser/extensions/blacklist.h"
+#include "chrome/browser/extensions/blacklist_fetcher.h"
 #include "chrome/browser/extensions/extension_prefs.h"
 #include "chrome/browser/extensions/fake_safe_browsing_database_manager.h"
 #include "chrome/browser/extensions/test_blacklist.h"
+#include "chrome/browser/extensions/test_blacklist_fetcher.h"
 #include "chrome/browser/extensions/test_extension_prefs.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace extensions {
 namespace {
 
 std::set<std::string> Set(const std::string& a) {
   std::set<std::string> set;
   set.insert(a);
   return set;
 }
 std::set<std::string> Set(const std::string& a, const std::string& b) {
   std::set<std::string> set = Set(a);
   set.insert(b);
   return set;
 }
 std::set<std::string> Set(const std::string& a,
                           const std::string& b,
                           const std::string& c) {
   std::set<std::string> set = Set(a, b);
   set.insert(c);
   return set;
 }
 std::set<std::string> Set(const std::string& a,
                           const std::string& b,
-                          const std::string& d,
-                          const std::string& c) {
+                          const std::string& c,
+                          const std::string& d) {
   std::set<std::string> set = Set(a, b, c);
   set.insert(d);
   return set;
 }
 
 class BlacklistTest : public testing::Test {
  public:
   BlacklistTest()
       : test_prefs_(base::MessageLoopProxy::current()),
         blacklist_db_(new FakeSafeBrowsingDatabaseManager(false)),
@@ skipping 15 matching lines @@
  private:
   content::TestBrowserThreadBundle browser_thread_bundle_;
 
   TestExtensionPrefs test_prefs_;
 
   scoped_refptr<FakeSafeBrowsingDatabaseManager> blacklist_db_;
 
   Blacklist::ScopedDatabaseManagerForTest scoped_blacklist_db_;
 };
 
-void Assign(std::set<std::string> *to, const std::set<std::string>& from) {
+class BlacklistFetcherMock : public BlacklistFetcher {
+ public:
+  virtual void Request(const std::string& id, const RequestCallback& callback) {
+    Blacklist::BlacklistState result = Blacklist::NOT_BLACKLISTED;
+    if (states_.find(id) != states_.end()) {

[not at google - send to devlin, 2013/11/18 17:03:49]

likewise you could use count here. but since you do actually access states_[id]
you could hold onto the result of find.

also nit, single-line if.

[Oleg Eterevsky, 2013/11/20 12:58:45]

If we were using C++11, I'd write
  (auto it = state_.find(it)) != states_.end()
Done.

+      result = states_[id];
+    }
+
+    base::MessageLoopProxy::current()->PostTask(FROM_HERE,
+                                                base::Bind(callback, result));
+  }
+
+  void SetState(const std::string& id, Blacklist::BlacklistState state) {
+    states_[id] = state;
+  }
+
+ private:
+  std::map<std::string, Blacklist::BlacklistState> states_;
+
+  bool delayed_;

[not at google - send to devlin, 2013/11/18 17:03:49]

not used

[Oleg Eterevsky, 2013/11/20 12:58:45]

Done.

+};
+
+void AssignSet(std::set<std::string> *to, const std::set<std::string>& from) {
   *to = from;
 }
 
+void AssignMap(Blacklist::BlacklistStateMap *to,
+            const Blacklist::BlacklistStateMap& from) {

[not at google - send to devlin, 2013/11/18 17:03:49]

nit: indent

though could you templatise Assign rather than making two methods?

[Oleg Eterevsky, 2013/11/20 12:58:45]

Done.

+  *to = from;
+}
+
+}  // namespace
+
 TEST_F(BlacklistTest, OnlyIncludesRequestedIDs) {
   std::string a = AddExtension("a");
   std::string b = AddExtension("b");
   std::string c = AddExtension("c");
 
   Blacklist blacklist(prefs());
   TestBlacklist tester(&blacklist);
+  BlacklistFetcherMock* fetcher = new BlacklistFetcherMock();
+  fetcher->SetState(a, Blacklist::BLACKLISTED_MALWARE);
+  fetcher->SetState(b, Blacklist::BLACKLISTED_MALWARE);
+  blacklist.SetBlacklistFetcherForTest(fetcher);
 
   blacklist_db()->Enable();
   blacklist_db()->SetUnsafe(a, b);
 
   EXPECT_EQ(Blacklist::BLACKLISTED_MALWARE, tester.GetBlacklistState(a));
   EXPECT_EQ(Blacklist::BLACKLISTED_MALWARE, tester.GetBlacklistState(b));
   EXPECT_EQ(Blacklist::NOT_BLACKLISTED, tester.GetBlacklistState(c));
 
   std::set<std::string> blacklisted_ids;
-  blacklist.GetMalwareIDs(Set(a, c), base::Bind(&Assign, &blacklisted_ids));
+  blacklist.GetMalwareIDs(Set(a, c), base::Bind(&AssignSet, &blacklisted_ids));
   base::RunLoop().RunUntilIdle();
 
   EXPECT_EQ(Set(a), blacklisted_ids);
 }
 
 TEST_F(BlacklistTest, SafeBrowsing) {
   std::string a = AddExtension("a");
 
   Blacklist blacklist(prefs());
   TestBlacklist tester(&blacklist);
+  BlacklistFetcherMock* fetcher = new BlacklistFetcherMock();
+  fetcher->SetState(a, Blacklist::BLACKLISTED_MALWARE);
+  blacklist.SetBlacklistFetcherForTest(fetcher);
 
   EXPECT_EQ(Blacklist::NOT_BLACKLISTED, tester.GetBlacklistState(a));
 
   blacklist_db()->SetUnsafe(a);
   // The manager is still disabled at this point, so it won't be blacklisted.
   EXPECT_EQ(Blacklist::NOT_BLACKLISTED, tester.GetBlacklistState(a));
 
   blacklist_db()->Enable().NotifyUpdate();
   base::RunLoop().RunUntilIdle();
   // Now it should be.
@@ skipping 15 matching lines @@
   // Blacklist some non-installed extensions. This is what the old preferences
   // blacklist looked like.
   std::string c = "cccccccccccccccccccccccccccccccc";
   std::string d = "dddddddddddddddddddddddddddddddd";
   prefs()->SetExtensionBlacklisted(c, true);
   prefs()->SetExtensionBlacklisted(d, true);
 
   EXPECT_EQ(Set(a, c, d), prefs()->GetBlacklistedExtensions());
 
   Blacklist blacklist(prefs());
-  TestBlacklist tester(&blacklist);
+  blacklist.SetBlacklistFetcherForTest(new BlacklistFetcherMock());
 
   // Blacklist has been cleared. Only the installed extension "a" left.
   EXPECT_EQ(Set(a), prefs()->GetBlacklistedExtensions());
   EXPECT_TRUE(prefs()->GetInstalledExtensionInfo(a).get());
   EXPECT_TRUE(prefs()->GetInstalledExtensionInfo(b).get());
 
   // "a" won't actually be *blacklisted* since it doesn't appear in
   // safebrowsing. Blacklist no longer reads from prefs. This is purely a
   // concern of somebody else (currently, ExtensionService).
   std::set<std::string> blacklisted_ids;
   blacklist.GetMalwareIDs(Set(a, b, c, d),
-                          base::Bind(&Assign, &blacklisted_ids));
+                          base::Bind(&AssignSet, &blacklisted_ids));
   base::RunLoop().RunUntilIdle();
   EXPECT_EQ(std::set<std::string>(), blacklisted_ids);
 
   // Prefs are still unaffected for installed extensions, though.
   EXPECT_TRUE(prefs()->IsExtensionBlacklisted(a));
   EXPECT_FALSE(prefs()->IsExtensionBlacklisted(b));
   EXPECT_FALSE(prefs()->IsExtensionBlacklisted(c));
   EXPECT_FALSE(prefs()->IsExtensionBlacklisted(d));
 }
 
-}  // namespace
+// Test getting different blacklist states from Blacklist.
+TEST_F(BlacklistTest, GetBlacklistStates) {
+  Blacklist blacklist(prefs());
+
+  std::string a = AddExtension("a");
+  std::string b = AddExtension("b");
+  std::string c = AddExtension("c");
+  std::string d = AddExtension("d");
+  std::string e = AddExtension("e");
+
+  blacklist_db()->Enable();
+  blacklist_db()->SetUnsafe(a, b, c, d);
+
+  BlacklistFetcherMock* fetcher = new BlacklistFetcherMock();
+  fetcher->SetState(a, Blacklist::BLACKLISTED_MALWARE);
+  fetcher->SetState(b, Blacklist::BLACKLISTED_SECURITY_VULNERABILITY);
+  fetcher->SetState(c, Blacklist::BLACKLISTED_CWS_POLICY_VIOLATION);
+  fetcher->SetState(d, Blacklist::BLACKLISTED_POTENTIALLY_UNWANTED);
+  blacklist.SetBlacklistFetcherForTest(fetcher);
+
+  Blacklist::BlacklistStateMap states_abc;
+  Blacklist::BlacklistStateMap states_bcd;
+  blacklist.GetBlacklistedIDs(Set(a, b, c, e),
+                              base::Bind(&AssignMap, &states_abc));
+  blacklist.GetBlacklistedIDs(Set(b, c, d, e),
+                              base::Bind(&AssignMap, &states_bcd));
+  base::RunLoop().RunUntilIdle();
+
+  EXPECT_EQ(Blacklist::BLACKLISTED_MALWARE, states_abc[a]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_SECURITY_VULNERABILITY, states_abc[b]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_CWS_POLICY_VIOLATION, states_abc[c]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_SECURITY_VULNERABILITY, states_bcd[b]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_CWS_POLICY_VIOLATION, states_bcd[c]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_POTENTIALLY_UNWANTED, states_bcd[d]);
+  EXPECT_EQ(states_abc.end(), states_abc.find(e));

[not at google - send to devlin, 2013/11/18 17:03:49]

colud use EXPECT_EQ(0, states_abc.count(e)) here.

[Oleg Eterevsky, 2013/11/20 12:58:45]

Same find vs count as before.

+  EXPECT_EQ(states_bcd.end(), states_bcd.find(e));

[not at google - send to devlin, 2013/11/18 17:03:49]

could you also test a purely cached response?

[Oleg Eterevsky, 2013/11/20 12:58:45]

Done.

+}
+
+// Test both Blacklist and BlacklistFetcher by requesting the blacklist states,
+// sending fake requests and parsing the responses.
+TEST_F(BlacklistTest, FetchBlacklistStates) {
+  Blacklist blacklist(prefs());
+
+  std::string a = AddExtension("a");
+  std::string b = AddExtension("b");
+  std::string c = AddExtension("c");
+
+  blacklist_db()->Enable();
+  blacklist_db()->SetUnsafe(a, b);
+
+  // Prepare real fetcher.
+  BlacklistFetcher* fetcher = new BlacklistFetcher();
+  TestBlacklistFetcher fetcher_tester(fetcher);
+  blacklist.SetBlacklistFetcherForTest(fetcher);
+
+  fetcher_tester.SetBlacklistVerdict(
+      a, ClientCRXListInfoResponse_Verdict_CWS_POLICY_VIOLATION);
+  fetcher_tester.SetBlacklistVerdict(
+      b, ClientCRXListInfoResponse_Verdict_POTENTIALLY_UNWANTED);
+
+  Blacklist::BlacklistStateMap states;
+  blacklist.GetBlacklistedIDs(Set(a, b, c), base::Bind(&AssignMap, &states));
+  base::RunLoop().RunUntilIdle();
+
+    // Two fetchers should be created.

[not at google - send to devlin, 2013/11/18 17:03:49]

indent

[Oleg Eterevsky, 2013/11/20 12:58:45]

Done.

+  EXPECT_TRUE(fetcher_tester.HandleFetcher(0));
+  EXPECT_TRUE(fetcher_tester.HandleFetcher(1));
+
+  EXPECT_EQ(Blacklist::BLACKLISTED_CWS_POLICY_VIOLATION, states[a]);
+  EXPECT_EQ(Blacklist::BLACKLISTED_POTENTIALLY_UNWANTED, states[b]);
+  EXPECT_EQ(states.end(), states.find(c));

[not at google - send to devlin, 2013/11/18 17:03:49]

likewise caching

[Oleg Eterevsky, 2013/11/20 12:58:45]

Done.

+}
+
 }  // namespace extensions