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 test safebrowsing database
+// and a test 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/protocol_manager.h"
+#include "chrome/browser/safe_browsing/safe_browsing_database.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"
+
+namespace {
+
+// A SafeBrowingDatabase class that allows us to inject the malicious URLs.
+class TestSafeBrowsingDatabase :  public SafeBrowsingDatabase {
+ public:
+  TestSafeBrowsingDatabase() {}
+
+  virtual ~TestSafeBrowsingDatabase() {}
+
+  // 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) {
+    base::hash_map<std::string, Hits>::const_iterator
+        badurls_it = badurls_.find(url.spec());
+    if (badurls_it == badurls_.end())
+      return false;
+    *prefix_hits = badurls_it->second.prefix_hits;
+    *full_hits = badurls_it->second.full_hits;
+    return true;
+  }
+
+  virtual bool UpdateStarted(std::vector<SBListChunkRanges>* lists) {
+    ADD_FAILURE() << "Not implemented.";
+    return false;
+  }
+  virtual void InsertChunks(const std::string& list_name,
+                            const SBChunkList& chunks) {
+    ADD_FAILURE() << "Not implemented.";
+  }
+  virtual void DeleteChunks(const std::vector<SBChunkDelete>& chunk_deletes) {
+    ADD_FAILURE() << "Not implemented.";
+  }
+  virtual void UpdateFinished(bool update_succeeded) {
+    ADD_FAILURE() << "Not implemented.";
+  }
+  virtual void CacheHashResults(const std::vector<SBPrefix>& prefixes,
+      const std::vector<SBFullHashResult>& full_hits) {
+    // Do nothing for the cache.
+    return;
+  }
+
+  // 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_;
+};
+
+// Factory that creates TestSafeBrowsingDatabase instances.
+class TestSafeBrowsingDatabaseFactory : public SafeBrowsingDatabaseFactory {
+ public:
+  TestSafeBrowsingDatabaseFactory() : db_(NULL) {}
+  virtual ~TestSafeBrowsingDatabaseFactory() {}
+
+  virtual SafeBrowsingDatabase* CreateSafeBrowsingDatabase() {
+    db_ = new TestSafeBrowsingDatabase();
+    return db_;
+  }
+  TestSafeBrowsingDatabase* GetDb() {
+    return db_;
+  }
+ private:
+  // Owned by the SafebrowsingService.
+  TestSafeBrowsingDatabase* db_;
+};
+
+
+// A TestProtocolManager that could return fixed responses from
+// safebrowsing server for testing purpose.
+class TestProtocolManager :  public SafeBrowsingProtocolManager {
+ public:
+  TestProtocolManager(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));
+  }
+
+  // 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 TestProtocolManager instances.
+class TestSBProtocolManagerFactory : public SBProtocolManagerFactory {
+ public:
+  TestSBProtocolManagerFactory() : pm_(NULL) {}
+  virtual ~TestSBProtocolManagerFactory() {}
+
+  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 TestProtocolManager(
+        sb_service, client_name, client_key, wrapped_key,
+        request_context_getter, info_url_prefix, mackey_url_prefix,
+        disable_auto_update);
+    return pm_;
+  }
+  TestProtocolManager* GetProtocolManager() {
+    return pm_;
+  }
+ private:
+  // Owned by the SafebrowsingService.
+  TestProtocolManager* 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() {
+    // InProcessBrowserTest::SetUp() intantiates SafebrowsingService and
+    // RegisterFactory has to be called before SafeBrowsingService is created.
+    SafeBrowsingDatabase::RegisterFactory(&db_factory_);
+    SafeBrowsingProtocolManager::RegisterFactory(&pm_factory_);
+
+    InProcessBrowserTest::SetUp();
+  }
+
+  virtual void TearDown() {
+    InProcessBrowserTest::TearDown();
+
+    // Unregister test factories after InProcessBrowserTest::TearDown
+    // (which destructs SafeBrowsingService).
+    SafeBrowsingDatabase::RegisterFactory(NULL);
+    SafeBrowsingProtocolManager::RegisterFactory(NULL);
+  }
+
+  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;
+    TestSafeBrowsingDatabase* db = db_factory_.GetDb();
+    db->AddUrl(url, prefix_hits, empty_full_hits);
+
+    TestProtocolManager* pm = pm_factory_.GetProtocolManager();
+    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_;
+  TestSBProtocolManagerFactory pm_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(SafeBrowsingServiceTest);
+};
+
+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