// Copyright 2014 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 <set>

#include "base/run_loop.h"
#include "base/strings/string_number_conversions.h"
#include "chrome/browser/extensions/extension_browsertest.h"
#include "chrome/browser/extensions/extension_service.h"
#include "chrome/browser/extensions/extension_storage_monitor.h"
#include "chrome/browser/ui/extensions/app_launch_params.h"
#include "chrome/browser/ui/extensions/application_launch.h"
#include "content/public/test/test_utils.h"
#include "extensions/browser/extension_prefs.h"
#include "extensions/browser/extension_registry.h"
#include "extensions/browser/extension_system.h"
#include "extensions/browser/test_extension_registry_observer.h"
#include "extensions/common/constants.h"
#include "extensions/test/extension_test_message_listener.h"
#include "ui/message_center/message_center.h"
#include "ui/message_center/message_center_observer.h"

namespace extensions {

namespace {

const int kInitialUsageThreshold = 500;

const char kWriteDataApp[] = "storage_monitor/write_data";

class NotificationObserver : public message_center::MessageCenterObserver {
 public:
  explicit NotificationObserver(const std::string& target_notification)
      : message_center_(message_center::MessageCenter::Get()),
        target_notification_id_(target_notification),
        waiting_(false) {
    message_center_->AddObserver(this);
  }

  ~NotificationObserver() override { message_center_->RemoveObserver(this); }

  bool HasReceivedNotification() const {
    return received_notifications_.find(target_notification_id_) !=
      received_notifications_.end();
  }

  // Runs the message loop and returns true if a notification is received.
  // Immediately returns true if a notification has already been received.
  bool WaitForNotification() {
    if (HasReceivedNotification())
      return true;

    waiting_ = true;
    content::RunMessageLoop();
    waiting_ = false;
    return HasReceivedNotification();
  }

 private:
  // MessageCenterObserver implementation:
  void OnNotificationAdded(const std::string& notification_id) override {
    received_notifications_.insert(notification_id);

    if (waiting_ && HasReceivedNotification())
      base::MessageLoopForUI::current()->Quit();
  }

  message_center::MessageCenter* message_center_;
  std::set<std::string> received_notifications_;
  std::string target_notification_id_;
  bool waiting_;
};

}  // namespace

class ExtensionStorageMonitorTest : public ExtensionBrowserTest {
 public:
  ExtensionStorageMonitorTest() : storage_monitor_(NULL) {}

 protected:
  // ExtensionBrowserTest overrides:
  void SetUpOnMainThread() override {
    ExtensionBrowserTest::SetUpOnMainThread();

    InitStorageMonitor();
  }

  ExtensionStorageMonitor* monitor() {
    CHECK(storage_monitor_);
    return storage_monitor_;
  }

  int64 GetInitialExtensionThreshold() {
    CHECK(storage_monitor_);
    return storage_monitor_->initial_extension_threshold_;
  }

  int64 GetInitialEphemeralThreshold() {
    CHECK(storage_monitor_);
    return storage_monitor_->initial_ephemeral_threshold_;
  }

  void DisableForInstalledExtensions() {
    CHECK(storage_monitor_);
    storage_monitor_->enable_for_all_extensions_ = false;
  }

  const Extension* InitWriteDataApp() {
    base::FilePath path = test_data_dir_.AppendASCII(kWriteDataApp);
    const Extension* extension = InstallExtension(path, 1);
    EXPECT_TRUE(extension);
    return extension;
  }

  const Extension* InitWriteDataEphemeralApp() {
    // The threshold for installed extensions should be higher than ephemeral
    // apps.
    storage_monitor_->initial_extension_threshold_ =
        storage_monitor_->initial_ephemeral_threshold_ * 4;

    base::FilePath path = test_data_dir_.AppendASCII(kWriteDataApp);
    const Extension* extension = InstallEphemeralAppWithSourceAndFlags(
        path, 1, Manifest::INTERNAL, Extension::NO_FLAGS);
    EXPECT_TRUE(extension);
    return extension;
  }

  std::string GetNotificationId(const std::string& extension_id) {
    return monitor()->GetNotificationId(extension_id);
  }

  bool IsStorageNotificationEnabled(const std::string& extension_id) {
    return monitor()->IsStorageNotificationEnabled(extension_id);
  }

  int64 GetNextStorageThreshold(const std::string& extension_id) {
    return monitor()->GetNextStorageThreshold(extension_id);
  }

  void WriteBytesExpectingNotification(const Extension* extension,
                                       int num_bytes) {
    int64 previous_threshold = GetNextStorageThreshold(extension->id());
    WriteBytes(extension, num_bytes, true);
    EXPECT_GT(GetNextStorageThreshold(extension->id()), previous_threshold);
  }

  void WriteBytesNotExpectingNotification(const Extension* extension,
                                         int num_bytes) {
    WriteBytes(extension, num_bytes, false);
  }

  void SimulateUninstallDialogAccept() {
    // Ensure the uninstall dialog was shown and fake an accept.
    ASSERT_TRUE(monitor()->uninstall_dialog_.get());
    monitor()->ExtensionUninstallAccepted();
  }

 private:
  void InitStorageMonitor() {
    storage_monitor_ = ExtensionStorageMonitor::Get(profile());
    ASSERT_TRUE(storage_monitor_);

    // Override thresholds so that we don't have to write a huge amount of data
    // to trigger notifications in these tests.
    storage_monitor_->enable_for_all_extensions_ = true;
    storage_monitor_->initial_extension_threshold_ = kInitialUsageThreshold;
    storage_monitor_->initial_ephemeral_threshold_ = kInitialUsageThreshold;

    // To ensure storage events are dispatched from QuotaManager immediately.
    storage_monitor_->observer_rate_ = base::TimeDelta();
  }

  // Write a number of bytes to persistent storage.
  void WriteBytes(const Extension* extension,
                  int num_bytes,
                  bool expected_notification) {
    ExtensionTestMessageListener launched_listener("launched", true);
    ExtensionTestMessageListener write_complete_listener(
        "write_complete", false);
    NotificationObserver notification_observer(
        GetNotificationId(extension->id()));

    OpenApplication(AppLaunchParams(profile(), extension, LAUNCH_CONTAINER_NONE,
                                    NEW_WINDOW, extensions::SOURCE_TEST));
    ASSERT_TRUE(launched_listener.WaitUntilSatisfied());

    // Instruct the app to write |num_bytes| of data.
    launched_listener.Reply(base::IntToString(num_bytes));
    ASSERT_TRUE(write_complete_listener.WaitUntilSatisfied());

    if (expected_notification) {
      EXPECT_TRUE(notification_observer.WaitForNotification());
    } else {
      base::RunLoop().RunUntilIdle();
      EXPECT_FALSE(notification_observer.HasReceivedNotification());
    }
  }

  ExtensionStorageMonitor* storage_monitor_;
};

// Control - No notifications should be shown if usage remains under the
// threshold.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, UnderThreshold) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesNotExpectingNotification(extension, 1);
}

// Ensure a notification is shown when usage reaches the first threshold.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, ExceedInitialThreshold) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension, GetInitialExtensionThreshold());
}

// Ensure a notification is shown when usage immediately exceeds double the
// first threshold.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, DoubleInitialThreshold) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension,
                                  GetInitialExtensionThreshold() * 2);
}

// Ensure that notifications are not fired if the next threshold has not been
// reached.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, ThrottleNotifications) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);

  // Exceed the first threshold.
  WriteBytesExpectingNotification(extension, GetInitialExtensionThreshold());

  // Stay within the next threshold.
  WriteBytesNotExpectingNotification(extension, 1);
}

// Verify that notifications are disabled when the user clicks the action button
// in the notification.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, UserDisabledNotifications) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension, GetInitialExtensionThreshold());

  EXPECT_TRUE(IsStorageNotificationEnabled(extension->id()));

  // Fake clicking the notification button to disable notifications.
  message_center::MessageCenter::Get()->ClickOnNotificationButton(
      GetNotificationId(extension->id()),
      ExtensionStorageMonitor::BUTTON_DISABLE_NOTIFICATION);

  EXPECT_FALSE(IsStorageNotificationEnabled(extension->id()));

  // Expect to receive no further notifications when usage continues to
  // increase.
  int64 next_threshold = GetNextStorageThreshold(extension->id());
  int64 next_data_size = next_threshold - GetInitialExtensionThreshold();
  ASSERT_GT(next_data_size, 0);

  WriteBytesNotExpectingNotification(extension, next_data_size);
}

// Verify that thresholds for ephemeral apps are reset when they are
// promoted to regular installed apps.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, EphemeralAppLowUsage) {
  const Extension* extension = InitWriteDataEphemeralApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension, GetInitialEphemeralThreshold());

  // Store the number of bytes until the next threshold is reached.
  int64 next_threshold = GetNextStorageThreshold(extension->id());
  int64 next_data_size = next_threshold - GetInitialEphemeralThreshold();
  ASSERT_GT(next_data_size, 0);
  EXPECT_GE(GetInitialExtensionThreshold(), next_threshold);

  // Promote the ephemeral app.
  ExtensionService* service =
      ExtensionSystem::Get(profile())->extension_service();
  service->PromoteEphemeralApp(extension, false);

  // The next threshold should now be equal to the initial threshold for
  // extensions (which is higher than the initial threshold for ephemeral apps).
  EXPECT_EQ(GetInitialExtensionThreshold(),
            GetNextStorageThreshold(extension->id()));

  // Since the threshold was increased, a notification should not be
  // triggered.
  WriteBytesNotExpectingNotification(extension, next_data_size);
}

// Verify that thresholds for ephemeral apps are not reset when they are
// promoted to regular installed apps if their usage is higher than the initial
// threshold for installed extensions.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, EphemeralAppWithHighUsage) {
  const Extension* extension = InitWriteDataEphemeralApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension, GetInitialExtensionThreshold());
  int64 saved_next_threshold = GetNextStorageThreshold(extension->id());

  // Promote the ephemeral app.
  ExtensionService* service =
      ExtensionSystem::Get(profile())->extension_service();
  service->PromoteEphemeralApp(extension, false);

  // The next threshold should not have changed.
  EXPECT_EQ(saved_next_threshold, GetNextStorageThreshold(extension->id()));
}

// Ensure that monitoring is disabled for installed extensions if
// |enable_for_all_extensions_| is false. This test can be removed if monitoring
// is eventually enabled for all extensions.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest,
                       DisableForInstalledExtensions) {
  DisableForInstalledExtensions();

  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesNotExpectingNotification(extension, GetInitialExtensionThreshold());
}

// Verify that notifications are disabled when the user clicks the action button
// in the notification.
IN_PROC_BROWSER_TEST_F(ExtensionStorageMonitorTest, UninstallExtension) {
  const Extension* extension = InitWriteDataApp();
  ASSERT_TRUE(extension);
  WriteBytesExpectingNotification(extension, GetInitialExtensionThreshold());

  // Fake clicking the notification button to uninstall.
  message_center::MessageCenter::Get()->ClickOnNotificationButton(
      GetNotificationId(extension->id()),
      ExtensionStorageMonitor::BUTTON_UNINSTALL);

  // Also fake accepting the uninstall.
  TestExtensionRegistryObserver observer(ExtensionRegistry::Get(profile()),
                                         extension->id());
  SimulateUninstallDialogAccept();
  observer.WaitForExtensionUninstalled();
}

}  // namespace extensions