Add tests for redirect responses from SafeBrowsingProtocolManager.

chrome/browser/safe_browsing/protocol_manager_unittest.cc

Index: chrome/browser/safe_browsing/protocol_manager_unittest.cc
diff --git a/chrome/browser/safe_browsing/protocol_manager_unittest.cc b/chrome/browser/safe_browsing/protocol_manager_unittest.cc
index 75f4097b89681bcfa7b62eedd9b55abfabcb2637..e969ae415a5771e41575ad0be2b2f29046457eaf 100644
--- a/chrome/browser/safe_browsing/protocol_manager_unittest.cc
+++ b/chrome/browser/safe_browsing/protocol_manager_unittest.cc
@@ -253,7 +253,8 @@ class MockProtocolDelegate : public SafeBrowsingProtocolManagerDelegate {
   MOCK_METHOD1(UpdateFinished, void(bool));
   MOCK_METHOD0(ResetDatabase, void());
   MOCK_METHOD1(GetChunks, void(GetChunksCallback));
-  MOCK_METHOD2(AddChunks, void(const std::string&, SBChunkList*));
+  MOCK_METHOD3(AddChunks, void(const std::string&, SBChunkList*,
+                               AddChunksCallback));
   MOCK_METHOD1(DeleteChunks, void(std::vector<SBChunkDelete>*));
 };
 
@@ -312,6 +313,15 @@ void InvokeGetChunksCallback(
   callback.Run(ranges, database_error);
 }
 
+void DeferAddChunksCallback(
+    SafeBrowsingProtocolManagerDelegate::AddChunksCallback callback) {
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner(
+      base::ThreadTaskRunnerHandle::Get());
+  if (!task_runner)
+    return;
+  task_runner->PostTask(FROM_HERE, callback);
+}
+
 }  // namespace
 
 // Tests that the Update protocol will be skipped if there are problems
@@ -603,3 +613,216 @@ TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseReset) {
   url_fetcher->SetResponseString("r:pleasereset\n");
   url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
 }
+
+// Tests a single valid update response, followed by a single redirect response
+// that has an valid, but empty body.
+TEST_F(SafeBrowsingProtocolManagerTest, EmptyRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);
+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString("");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+}
+
+// Tests a single valid update response, followed by a single redirect response
+// that has an invalid body.
+TEST_F(SafeBrowsingProtocolManagerTest, InvalidRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);
+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString("THIS IS AN INVALID RESPONSE");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+}
+
+// Tests a single valid update response, followed by a single redirect response
+// containing chunks.
+TEST_F(SafeBrowsingProtocolManagerTest, ChunksRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, AddChunks("goog-phish-shavar", _, _)).WillOnce(
+      testing::WithArgs<2>(
+          Invoke(DeferAddChunksCallback)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);
+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString(
+      "a:4:4:9\n"
+      "host\1fdaf");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+
+  // The AddChunksCallback needs to be invoked
+  runner->RunTasks();
+}
+
+// Tests a single valid update response, followed by multiple redirect responses
+// containing chunks.
+TEST_F(SafeBrowsingProtocolManagerTest, ChunksMultipleRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, AddChunks("goog-phish-shavar", _, _)).
+      WillRepeatedly(testing::WithArgs<2>(
+          Invoke(DeferAddChunksCallback)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/one\n"
+      "u:redirect-server.example.com/two\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // First redirect request.
+  net::TestURLFetcher* first_chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(first_chunk_url_fetcher != NULL);
+  first_chunk_url_fetcher->set_status(net::URLRequestStatus());
+  first_chunk_url_fetcher->set_response_code(200);
+  first_chunk_url_fetcher->SetResponseString(
+      "a:4:4:9\n"
+      "host\1aaaa");
+  first_chunk_url_fetcher->delegate()->OnURLFetchComplete(
+      first_chunk_url_fetcher);
+
+  // Invoke the AddChunksCallback
+  runner->RunTasks();
+
+  // Second redirect request.
+  net::TestURLFetcher* second_chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(2);
+  ASSERT_TRUE(second_chunk_url_fetcher != NULL);
+  second_chunk_url_fetcher->set_status(net::URLRequestStatus());
+  second_chunk_url_fetcher->set_response_code(200);
+  second_chunk_url_fetcher->SetResponseString(
+      "a:5:4:9\n"
+      "host\1bbbb");
+  second_chunk_url_fetcher->delegate()->OnURLFetchComplete(
+      second_chunk_url_fetcher);
+
+  // The AddChunksCallback needs to be invoked
+  runner->RunTasks();
+}