Implement FCM handler for WebProtect binary uploads

This class will handle the interactions with FCM, routing messages to the
appropriate callback, and managing the lifetime of the InstanceID token.
At the moment, the InstanceID token is created and never deleted, but in
the future we may want to manage the lifetime of the token.

Bug: 980784
Change-Id: I605affb03cd2e01eb258b15ebdc147b15a022a3c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1731542
Commit-Queue: Daniel Rubery <drubery@chromium.org>
Reviewed-by: Varun Khaneja <vakh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#685231}

chrome/browser/safe_browsing/BUILD.gn

@@ -128,6 +128,8 @@ jumbo_static_library("safe_browsing") {
         "client_side_detection_service.h",
         "client_side_model_loader.cc",
         "client_side_model_loader.h",
+        "download_protection/binary_fcm_service.cc",
+        "download_protection/binary_fcm_service.h",
         "download_protection/binary_upload_service.cc",
         "download_protection/binary_upload_service.h",
         "download_protection/check_client_download_request.cc",

chrome/browser/safe_browsing/download_protection/binary_fcm_service.cc

+// Copyright 2019 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/download_protection/binary_fcm_service.h"
+
+#include <memory>
+
+#include "base/logging.h"
+#include "chrome/browser/gcm/gcm_profile_service_factory.h"
+#include "chrome/browser/gcm/instance_id/instance_id_profile_service_factory.h"
+#include "chrome/browser/profiles/profile.h"
+#include "components/gcm_driver/gcm_driver.h"
+#include "components/gcm_driver/gcm_profile_service.h"
+#include "components/gcm_driver/instance_id/instance_id_driver.h"
+#include "components/gcm_driver/instance_id/instance_id_profile_service.h"
+#include "components/safe_browsing/proto/webprotect.pb.h"
+
+namespace safe_browsing {
+
+namespace {
+
+const char kBinaryFCMServiceAppId[] = "safe_browsing_fcm_service";
+// TODO(drubery): Once the server side has finalized their sender id, fill this
+// in.
+const char kBinaryFCMServiceSenderId[] = "SenderID";
+const char kBinaryFCMServiceMessageKey[] = "proto";
+
+}  // namespace
+
+const char BinaryFCMService::kInvalidId[] = "";
+
+// static
+std::unique_ptr<BinaryFCMService> BinaryFCMService::Create(Profile* profile) {
+  gcm::GCMProfileService* gcm_profile_service =
+      gcm::GCMProfileServiceFactory::GetForProfile(profile);
+  if (!gcm_profile_service)
+    return nullptr;
+
+  gcm::GCMDriver* gcm_driver = gcm_profile_service->driver();
+  if (!gcm_driver)
+    return nullptr;
+
+  instance_id::InstanceIDProfileService* instance_id_profile_service =
+      instance_id::InstanceIDProfileServiceFactory::GetForProfile(profile);
+  if (!instance_id_profile_service)
+    return nullptr;
+
+  instance_id::InstanceIDDriver* instance_id_driver =
+      instance_id_profile_service->driver();
+  if (!instance_id_driver)
+    return nullptr;
+
+  return std::make_unique<BinaryFCMService>(gcm_driver, instance_id_driver);
+}
+
+BinaryFCMService::BinaryFCMService(
+    gcm::GCMDriver* gcm_driver,
+    instance_id::InstanceIDDriver* instance_id_driver)
+    : gcm_driver_(gcm_driver),
+      instance_id_(kInvalidId),
+      weakptr_factory_(this) {
+  gcm_driver->AddAppHandler(kBinaryFCMServiceAppId, this);
+  instance_id_driver->GetInstanceID(kBinaryFCMServiceAppId)
+      ->GetToken(kBinaryFCMServiceSenderId, instance_id::kGCMScope,
+                 /*options=*/{},
+                 /*flags=*/{},
+                 base::BindOnce(&BinaryFCMService::OnGetInstanceID,
+                                weakptr_factory_.GetWeakPtr()));
+}
+
+BinaryFCMService::~BinaryFCMService() {
+  gcm_driver_->RemoveAppHandler(kBinaryFCMServiceAppId);
+}
+
+void BinaryFCMService::GetInstanceID(GetInstanceIDCallback callback) {
+  std::move(callback).Run(instance_id_);
+}
+
+void BinaryFCMService::SetCallbackForToken(
+    const std::string& token,
+    base::RepeatingCallback<void(DeepScanningClientResponse)> callback) {
+  message_token_map_[token] = std::move(callback);
+}
+
+void BinaryFCMService::ClearCallbackForToken(const std::string& token) {
+  message_token_map_.erase(token);
+}
+
+void BinaryFCMService::OnGetInstanceID(const std::string& instance_id,
+                                       instance_id::InstanceID::Result result) {
+  if (result == instance_id::InstanceID::Result::SUCCESS) {
+    instance_id_ = instance_id;
+  }
+}
+
+void BinaryFCMService::ShutdownHandler() {
+  NOTIMPLEMENTED();
+}
+
+void BinaryFCMService::OnStoreReset() {
+  NOTIMPLEMENTED();
+}
+
+void BinaryFCMService::OnMessage(const std::string& app_id,
+                                 const gcm::IncomingMessage& message) {
+  auto serialized_proto = message.data.find(kBinaryFCMServiceMessageKey);
+  if (serialized_proto == message.data.end())
+    return;
+
+  DeepScanningClientResponse response;
+  if (!response.ParseFromString(serialized_proto->second))
+    return;
+
+  auto callback_it = message_token_map_.find(response.token());
+  if (callback_it == message_token_map_.end())
+    return;
+
+  callback_it->second.Run(std::move(response));
+}
+
+void BinaryFCMService::OnMessagesDeleted(const std::string& app_id) {
+  NOTIMPLEMENTED();
+}
+
+void BinaryFCMService::OnSendError(
+    const std::string& app_id,
+    const gcm::GCMClient::SendErrorDetails& send_error_details) {
+  NOTIMPLEMENTED();
+}
+
+void BinaryFCMService::OnSendAcknowledged(const std::string& app_id,
+                                          const std::string& message_id) {
+  NOTIMPLEMENTED();
+}
+
+bool BinaryFCMService::CanHandle(const std::string& app_id) const {
+  return app_id == kBinaryFCMServiceAppId;
+}
+
+}  // namespace safe_browsing

chrome/browser/safe_browsing/download_protection/binary_fcm_service.h

+// Copyright 2019 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.
+
+#ifndef CHROME_BROWSER_SAFE_BROWSING_DOWNLOAD_PROTECTION_BINARY_FCM_SERVICE_H_
+#define CHROME_BROWSER_SAFE_BROWSING_DOWNLOAD_PROTECTION_BINARY_FCM_SERVICE_H_
+
+#include "base/memory/weak_ptr.h"
+#include "components/gcm_driver/gcm_app_handler.h"
+#include "components/gcm_driver/instance_id/instance_id.h"
+#include "components/safe_browsing/proto/webprotect.pb.h"
+
+class Profile;
+
+namespace gcm {
+class GCMDriver;
+}
+
+namespace instance_id {
+class InstanceIDDriver;
+}
+
+namespace safe_browsing {
+
+// This class handles the interactions with FCM for asynchronously receiving
+// notifications when scans are complete. It has two major responsibilities:
+// management of the InstanceID tokens, and routing of incoming FCM messages
+// to the appropriate callbacks.
+// This class is neither copyable nor movable.
+class BinaryFCMService : public gcm::GCMAppHandler {
+ public:
+  static std::unique_ptr<BinaryFCMService> Create(Profile* profile);
+
+  BinaryFCMService(gcm::GCMDriver* gcm_driver,
+                   instance_id::InstanceIDDriver* instance_id_driver);
+  ~BinaryFCMService() override;
+  BinaryFCMService(BinaryFCMService&) = delete;
+  BinaryFCMService& operator=(BinaryFCMService&) = delete;
+  BinaryFCMService(BinaryFCMService&&) = delete;
+  BinaryFCMService& operator=(BinaryFCMService&&) = delete;
+
+  using GetInstanceIDCallback =
+      base::OnceCallback<void(const std::string& token)>;
+  using OnMessageCallback =
+      base::RepeatingCallback<void(DeepScanningClientResponse)>;
+
+  void GetInstanceID(GetInstanceIDCallback callback);
+  void SetCallbackForToken(const std::string& token,
+                           OnMessageCallback callback);
+  void ClearCallbackForToken(const std::string& token);
+
+  // GCMAppHandler implementation
+  void ShutdownHandler() override;
+  void OnStoreReset() override;
+  void OnMessage(const std::string& app_id,
+                 const gcm::IncomingMessage& message) override;
+  void OnMessagesDeleted(const std::string& app_id) override;
+  void OnSendError(
+      const std::string& app_id,
+      const gcm::GCMClient::SendErrorDetails& send_error_details) override;
+  void OnSendAcknowledged(const std::string& app_id,
+                          const std::string& message_id) override;
+  bool CanHandle(const std::string& app_id) const override;
+
+  static const char kInvalidId[];
+
+ private:
+  void OnGetInstanceID(const std::string& instance_id,
+                       instance_id::InstanceID::Result result);
+
+  gcm::GCMDriver* gcm_driver_;
+  std::string instance_id_;
+  std::unordered_map<std::string, OnMessageCallback> message_token_map_;
+
+  base::WeakPtrFactory<BinaryFCMService> weakptr_factory_;
+};
+
+}  // namespace safe_browsing
+
+#endif  // CHROME_BROWSER_SAFE_BROWSING_DOWNLOAD_PROTECTION_BINARY_FCM_SERVICE_H_

chrome/browser/safe_browsing/download_protection/binary_fcm_service_unittest.cc

+// Copyright 2019 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/download_protection/binary_fcm_service.h"
+
+#include "base/run_loop.h"
+#include "chrome/browser/gcm/gcm_profile_service_factory.h"
+#include "chrome/test/base/testing_profile.h"
+#include "components/gcm_driver/common/gcm_message.h"
+#include "components/gcm_driver/fake_gcm_profile_service.h"
+#include "components/safe_browsing/proto/webprotect.pb.h"
+#include "content/public/test/test_browser_thread_bundle.h"
+#include "content/public/test/test_utils.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace safe_browsing {
+
+namespace {
+
+std::unique_ptr<KeyedService> BuildFakeGCMProfileService(
+    content::BrowserContext* context) {
+  return gcm::FakeGCMProfileService::Build(static_cast<Profile*>(context));
+}
+
+}  // namespace
+
+class BinaryFCMServiceTest : public ::testing::Test {
+ public:
+  BinaryFCMServiceTest() {
+    gcm::GCMProfileServiceFactory::GetInstance()->SetTestingFactory(
+        &profile_, base::BindRepeating(&BuildFakeGCMProfileService));
+
+    binary_fcm_service_ = BinaryFCMService::Create(&profile_);
+  }
+
+  Profile* profile() { return &profile_; }
+
+ protected:
+  content::TestBrowserThreadBundle thread_bundle_;
+  TestingProfile profile_;
+  std::unique_ptr<BinaryFCMService> binary_fcm_service_;
+};
+
+TEST_F(BinaryFCMServiceTest, GetsInstanceID) {
+  std::string received_instance_id = BinaryFCMService::kInvalidId;
+
+  // Allow |binary_fcm_service_| to get an instance id.
+  content::RunAllTasksUntilIdle();
+
+  binary_fcm_service_->GetInstanceID(base::BindOnce(
+      [](std::string* target_id, const std::string& instance_id) {
+        *target_id = instance_id;
+      },
+      &received_instance_id));
+
+  content::RunAllTasksUntilIdle();
+
+  EXPECT_NE(received_instance_id, BinaryFCMService::kInvalidId);
+}
+
+TEST_F(BinaryFCMServiceTest, RoutesMessages) {
+  DeepScanningClientResponse response1;
+  DeepScanningClientResponse response2;
+
+  binary_fcm_service_->SetCallbackForToken(
+      "token1", base::BindRepeating(
+                    [](DeepScanningClientResponse* target_response,
+                       DeepScanningClientResponse response) {
+                      *target_response = response;
+                    },
+                    &response1));
+  binary_fcm_service_->SetCallbackForToken(
+      "token2", base::BindRepeating(
+                    [](DeepScanningClientResponse* target_response,
+                       DeepScanningClientResponse response) {
+                      *target_response = response;
+                    },
+                    &response2));
+
+  DeepScanningClientResponse message;
+  std::string serialized_message;
+  gcm::IncomingMessage incoming_message;
+
+  // Test that a message with token1 is routed only to the first callback.
+  message.set_token("token1");
+  ASSERT_TRUE(message.SerializeToString(&serialized_message));
+  incoming_message.data["proto"] = serialized_message;
+  binary_fcm_service_->OnMessage("app_id", incoming_message);
+  EXPECT_EQ(response1.token(), "token1");
+  EXPECT_EQ(response2.token(), "");
+
+  // Test that a message with token2 is routed only to the second callback.
+  message.set_token("token2");
+  ASSERT_TRUE(message.SerializeToString(&serialized_message));
+  incoming_message.data["proto"] = serialized_message;
+  binary_fcm_service_->OnMessage("app_id", incoming_message);
+  EXPECT_EQ(response1.token(), "token1");
+  EXPECT_EQ(response2.token(), "token2");
+
+  // Test that I can clear a callback
+  response2.clear_token();
+  binary_fcm_service_->ClearCallbackForToken("token2");
+  binary_fcm_service_->OnMessage("app_id", incoming_message);
+  EXPECT_EQ(response2.token(), "");
+}
+
+}  // namespace safe_browsing

chrome/test/BUILD.gn

@@ -4522,6 +4522,7 @@ test("unit_tests") {
       "../browser/safe_browsing/client_side_detection_host_unittest.cc",
       "../browser/safe_browsing/client_side_detection_service_unittest.cc",
       "../browser/safe_browsing/client_side_model_loader_unittest.cc",
+      "../browser/safe_browsing/download_protection/binary_fcm_service_unittest.cc",
       "../browser/safe_browsing/download_protection/download_feedback_service_unittest.cc",
       "../browser/safe_browsing/download_protection/download_feedback_unittest.cc",
       "../browser/safe_browsing/download_protection/download_protection_service_unittest.cc",