Add a browser test for safebrowsing service....

chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc

+// Copyright (c) 2010 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.
+//
+// This test creates a safebrowsing service using fake safebrowsing database
+// and a fake protocol manager. It is used to test logics in safebrowsing
+// service.
+
+#include "base/command_line.h"
+#include "base/sha2.h"
+#include "chrome/browser/browser_process.h"
+#include "chrome/browser/browser_thread.h"
+#include "chrome/browser/safe_browsing/safe_browsing_database.h"
+#include "chrome/browser/safe_browsing/protocol_manager.h"
+#include "chrome/browser/safe_browsing/safe_browsing_service.h"
+#include "chrome/browser/safe_browsing/safe_browsing_util.h"
+#include "chrome/browser/tab_contents/tab_contents.h"
+#include "chrome/browser/tab_contents/tab_contents_view.h"
+#include "chrome/browser/ui/browser.h"
+#include "chrome/common/chrome_switches.h"
+#include "chrome/test/in_process_browser_test.h"
+#include "chrome/test/ui_test_utils.h"
+
+// A SafeBrowingDatabase class that allows us to inject the malicious URLs.
+class FakeSafeBrowsingDatabase :  public SafeBrowsingDatabase {
+ public:
+  FakeSafeBrowsingDatabase() {}
+
+  virtual ~FakeSafeBrowsingDatabase() {}
+
+  // Initializes the database with the given filename.
+  virtual void Init(const FilePath& filename) {}
+
+  // Deletes the current database and creates a new one.
+  virtual bool ResetDatabase() {
+    badurls_.clear();
+    return true;
+  }
+
+  // Called on the IO thread to check if the given URL is safe or not.  If we
+  // can synchronously determine that the URL is safe, CheckUrl returns true,
+  // otherwise it returns false.
+  virtual bool ContainsUrl(const GURL& url,
+                           std::string* matching_list,
+                           std::vector<SBPrefix>* prefix_hits,
+                           std::vector<SBFullHashResult>* full_hits,
+                           base::Time last_update) {
+    if (badurls_.find(url.spec()) == badurls_.end())

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

I'd be happy enough with badurls_.count(), here, but given that these all have
to be consistent to make sense, I think it would be reasonable to capture the
iterator in a variable and use that for the rest of the function.

[lzheng, 2010/12/08 01:57:32]

Done.

+      return false;
+    *prefix_hits = badurls_[url.spec()].prefix_hits;
+    *full_hits = badurls_[url.spec()].full_hits;
+    return true;
+  }
+
+  virtual bool UpdateStarted(std::vector<SBListChunkRanges>* lists) {
+    NOTREACHED();  // TODO(lzheng): Implement this for future tests.

[Paweł Hajdan Jr., 2010/12/07 08:17:22]

I suggest removing those NOTREACHED()s unless you like crashy tests. ADD_FAILURE
sounds like good replacement.

[lzheng, 2010/12/08 01:57:32]

Done.

+    return false;
+  }
+  virtual void InsertChunks(const std::string& list_name,
+                            const SBChunkList& chunks) {
+    NOTREACHED();  // TODO(lzheng): Implement this for future tests.
+  }
+  virtual void DeleteChunks(const std::vector<SBChunkDelete>& chunk_deletes) {
+    NOTREACHED();  // TODO(lzheng): Implement this for future tests.
+  }
+  virtual void UpdateFinished(bool update_succeeded) {
+    NOTREACHED();  // TODO(lzheng): Implement this for future tests.
+  }
+  virtual void CacheHashResults(const std::vector<SBPrefix>& prefixes,
+      const std::vector<SBFullHashResult>& full_hits) {
+    //    NOTREACHED();  // TODO(lzheng): Implement this for future tests.
+  }
+
+  // Fill up the database with test URL.
+  void AddUrl(const GURL& url,
+              const std::vector<SBPrefix>& prefix_hits,
+              const std::vector<SBFullHashResult>& full_hits) {
+    badurls_[url.spec()].prefix_hits = prefix_hits;
+    badurls_[url.spec()].full_hits = full_hits;
+  }
+ private:
+  struct Hits {
+    std::vector<SBPrefix> prefix_hits;
+    std::vector<SBFullHashResult> full_hits;
+  };
+  base::hash_map<std::string, Hits> badurls_;

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Why not a hash_multimap<>?  [Answer: Because that would be a bad idea :-).]

+};
+
+// Factory that creates FakeSafeBrowsingDatabase instances.
+class TestSafeBrowsingDatabaseFactory : public SafeBrowsingDatabaseFactory {

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Would rather you used consistent naming.  Fake* or Test* or Mock* for both the
object and the factory.

[lzheng, 2010/12/08 01:57:32]

"Test " then.

On 2010/12/07 21:22:57, shess wrote:

[lzheng, 2010/12/08 01:57:32]

"Test " then.

On 2010/12/07 21:22:57, shess wrote:

+ public:
+  TestSafeBrowsingDatabaseFactory() : db_(NULL) { }
+  virtual ~TestSafeBrowsingDatabaseFactory() { }
+
+  virtual SafeBrowsingDatabase* CreateSafeBrowsingDatabase() {
+    db_ = new FakeSafeBrowsingDatabase();
+    return db_;
+  }
+  SafeBrowsingDatabase* GetDb() {
+    return db_;
+  }
+ private:
+  // Owned by the Safebrowsing service.
+  SafeBrowsingDatabase* db_;
+};
+
+
+// A FakeProtocolManager that could return fixed responses from
+// safebrowsing server for testing purpose.
+class FakeProtocolManager :  public SafeBrowsingProtocolManager {
+ public:
+  FakeProtocolManager(SafeBrowsingService* sb_service,
+                      const std::string& client_name,
+                      const std::string& client_key,
+                      const std::string& wrapped_key,
+                      URLRequestContextGetter* request_context_getter,
+                      const std::string& info_url_prefix,
+                      const std::string& mackey_url_prefix,
+                      bool disable_auto_update)
+      : SafeBrowsingProtocolManager(sb_service, client_name, client_key,
+                                    wrapped_key, request_context_getter,
+                                    info_url_prefix, mackey_url_prefix,
+                                    disable_auto_update),
+        sb_service_(sb_service) {
+  }
+
+  // This function is called when there is a prefix hit in local safebrowsing
+  // database and safebrowsing service issues a get hash request to backends.
+  // We return a result from the prefilled full_hashes_ hash_map to simulate
+  // server's response.
+  virtual void GetFullHash(SafeBrowsingService::SafeBrowsingCheck* check,
+                           const std::vector<SBPrefix>& prefixes) {
+    // The hash result should be inserted to the full_hashes_.
+    ASSERT_TRUE(full_hashes_.find(check->url.spec()) != full_hashes_.end());
+    // When we get a valid response, always cache the result.
+    bool cancache = true;
+    BrowserThread::PostTask(
+        BrowserThread::IO, FROM_HERE,
+        NewRunnableMethod(
+            sb_service_, &SafeBrowsingService::HandleGetHashResults,
+            check, full_hashes_[check->url.spec()], cancache));

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

This scares me a little.  I can live with it.

+  }
+
+  // Prepare the GetFullHash results for |url|.
+  void SetGetFullHashResponse(const GURL& url,
+                              const SBFullHashResult& full_hash_result) {
+    full_hashes_[url.spec()].push_back(full_hash_result);
+  }
+
+ private:
+  base::hash_map<std::string, std::vector<SBFullHashResult> > full_hashes_;
+  SafeBrowsingService* sb_service_;
+};
+
+// Factory that creates FakeSafeBrowsingDatabase instances.

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Wrong class name in comment.

Again consistent naming between test object and text-object factory.

[lzheng, 2010/12/08 01:57:32]

Done.

+class TestProtocolManagerFactory : public ProtocolManagerFactory {
+ public:
+  TestProtocolManagerFactory() : pm_(NULL) { }
+  virtual ~TestProtocolManagerFactory() { }
+
+  virtual SafeBrowsingProtocolManager* CreateProtocolManager(
+      SafeBrowsingService* sb_service,
+      const std::string& client_name,
+      const std::string& client_key,
+      const std::string& wrapped_key,
+      URLRequestContextGetter* request_context_getter,
+      const std::string& info_url_prefix,
+      const std::string& mackey_url_prefix,
+      bool disable_auto_update) {
+    pm_ = new FakeProtocolManager(
+        sb_service, client_name, client_key, wrapped_key,
+        request_context_getter, info_url_prefix, mackey_url_prefix,
+        disable_auto_update);
+    return pm_;
+  }
+  SafeBrowsingProtocolManager* GetProtocolManager() {
+    return pm_;
+  }
+ private:
+  // Owned by the Safebrowsing service.
+  SafeBrowsingProtocolManager* pm_;
+};
+
+
+// Tests the safe browsing blocking page in a browser.
+class SafeBrowsingServiceTest : public InProcessBrowserTest {
+ public:
+  SafeBrowsingServiceTest() {
+  }
+
+  static void GenerateFullhashResult(const GURL& url,
+                                     const std::string& list_name,
+                                     int add_chunk_id,
+                                     SBFullHashResult* full_hash) {
+    std::string host;
+    std::string path;
+    safe_browsing_util::CanonicalizeUrl(url, &host, &path, NULL);
+    base::SHA256HashString(host + path, &full_hash->hash,
+                           sizeof(SBFullHash));
+    full_hash->list_name = list_name;
+    full_hash->add_chunk_id = add_chunk_id;
+  }
+
+  virtual void SetUp() {
+    SafeBrowsingDatabase::RegisterFactory(&db_factory_);
+    SafeBrowsingProtocolManager::RegisterFactory(&pm_factory_);
+    InProcessBrowserTest::SetUp();
+  }
+
+  virtual void TearDown() {
+    InProcessBrowserTest::TearDown();
+    SafeBrowsingDatabase::RegisterFactory(NULL);

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Un-register both factories.

I suspect you have the registration intentionally outside the nested
SetUp()/TearDown(), so that things aren't already setup.  If so, I think it's
subtle enough to comment on.

[lzheng, 2010/12/08 01:57:32]

Done.

+  }
+
+  virtual void SetUpCommandLine(CommandLine* command_line) {
+    // Makes sure the auto update is not triggered during the test.
+    // This test will fill up the database using testing prefixes
+    // and urls.
+    command_line->AppendSwitch(switches::kSbDisableAutoUpdate);
+  }
+
+  virtual void SetUpInProcessBrowserTestFixture() {
+    ASSERT_TRUE(test_server()->Start());
+  }
+
+  // This will setup the prefix in database and prepare protocol manager
+  // to response with |full_hash| for get full hash request.
+  void SetupResponseForUrl(const GURL& url,
+                            const SBFullHashResult& full_hash) {
+    std::vector<SBPrefix> prefix_hits;
+    prefix_hits.push_back(full_hash.hash.prefix);
+
+    // Make sure the full hits is empty unless we need to test the
+    // full hash is hit in database's local cache.
+    std::vector<SBFullHashResult> empty_full_hits;
+    FakeSafeBrowsingDatabase* db =
+        static_cast<FakeSafeBrowsingDatabase*>(db_factory_.GetDb());

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Just have db_factory_.GetDb() return FakeSafeBrowsingDatabase directly.  Though
the parent class has SafeBrowsingDatabase, I think you can return a proper
subclass.

Otherwise, just add another accessor, or a friend.  static_cast<>() shouldn't be
necessary when you have the fully-specified factory.

[lzheng, 2010/12/08 01:57:32]

Done.

+    db->AddUrl(url, prefix_hits, empty_full_hits);
+
+    FakeProtocolManager* pm =
+        static_cast<FakeProtocolManager*>(pm_factory_.GetProtocolManager());

[Scott Hess - ex-Googler, 2010/12/07 21:22:57]

Here too.

[lzheng, 2010/12/08 01:57:32]

Done.

+    pm->SetGetFullHashResponse(url, full_hash);
+  }
+
+  bool ShowingInterstitialPage() {
+    TabContents* contents = browser()->GetSelectedTabContents();
+    InterstitialPage* interstitial_page = contents->interstitial_page();
+    return interstitial_page != NULL;
+  }
+
+ private:
+  TestSafeBrowsingDatabaseFactory db_factory_;
+  TestProtocolManagerFactory pm_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(SafeBrowsingServiceTest);
+};
+
+namespace {

[Paweł Hajdan Jr., 2010/12/07 08:17:22]

nit: Is it possible to put everything here in the anonymous namespace?

[lzheng, 2010/12/08 01:57:32]

Done.

+
+const char kEmptyPage[] = "files/empty.html";
+const char kMalwarePage[] = "files/safe_browsing/malware.html";
+const char kMalwareIframe[] = "files/safe_browsing/malware_iframe.html";
+
+IN_PROC_BROWSER_TEST_F(SafeBrowsingServiceTest, Malware) {
+  GURL url = test_server()->GetURL(kEmptyPage);
+  // After adding the url to safebrowsing database and getfullhash result,
+  // we should see the interstitial page.
+  SBFullHashResult malware_full_hash;
+  int chunk_id = 0;
+  GenerateFullhashResult(url, safe_browsing_util::kMalwareList, chunk_id,
+                         &malware_full_hash);
+  SetupResponseForUrl(url, malware_full_hash);
+  ui_test_utils::NavigateToURL(browser(), url);
+  EXPECT_TRUE(ShowingInterstitialPage());
+}
+
+}  // namespace