[Extensions] Refactor StorageFrontend to own GetBytesInUse implementation.

Move the implementation of GetBytesInUse to StorageFrontend
and call it from storage_api.cc.

This is a part of larger work to move the extension storage
implementation to StorageFrontend, allowing storage to be
managed from outside of extension API calls and therefore
to expose it in other contexts (like CDP).

Bug: 40963428
Change-Id: If214d93911dc572b7721f5dffaa57080902b8ae6
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5463137
Reviewed-by: Devlin Cronin <rdevlin.cronin@chromium.org>
Commit-Queue: Oliver Dunk <oliverdunk@chromium.org>
Reviewed-by: Emilia Paz <emiliapaz@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1301233}

extensions/browser/api/storage/storage_api.cc

@@ -147,16 +147,21 @@ bool SettingsFunction::ShouldSkipQuotaLimiting() const {
          StorageAreaNamespace::kSync;
 }
 
-ExtensionFunction::ResponseAction SettingsFunction::Run() {
-  EXTENSION_FUNCTION_VALIDATE(args().size() >= 1);
-  EXTENSION_FUNCTION_VALIDATE(args()[0].is_string());
+bool SettingsFunction::PreRunValidation(std::string* error) {
+  if (!ExtensionFunction::PreRunValidation(error)) {
+    return false;
+  }
+
+  EXTENSION_FUNCTION_PRERUN_VALIDATE(args().size() >= 1);
+  EXTENSION_FUNCTION_PRERUN_VALIDATE(args()[0].is_string());
 
   // Not a ref since we remove the underlying value after.
   std::string storage_area_string = args()[0].GetString();
 
   mutable_args().erase(args().begin());
   storage_area_ = StorageAreaFromString(storage_area_string);
-  EXTENSION_FUNCTION_VALIDATE(storage_area_ != StorageAreaNamespace::kInvalid);
+  EXTENSION_FUNCTION_PRERUN_VALIDATE(storage_area_ !=
+                                     StorageAreaNamespace::kInvalid);
 
   // Session is the only storage area that does not use ValueStore, and will
   // return synchronously.
@@ -164,34 +169,46 @@ ExtensionFunction::ResponseAction SettingsFunction::Run() {
     // Currently only `session` can restrict the storage access. This call will
     // be moved after the other storage areas allow it.
     if (!IsAccessToStorageAllowed()) {
-      return RespondNow(
-          Error("Access to storage is not allowed from this context."));
+      *error = "Access to storage is not allowed from this context.";
+      return false;
     }
-    return RespondNow(RunInSession());
+    return true;
   }
 
   // All other StorageAreas use ValueStore with settings_namespace, and will
   // return asynchronously if successful.
   settings_namespace_ = StorageAreaToSettingsNamespace(storage_area_);
-  EXTENSION_FUNCTION_VALIDATE(settings_namespace_ !=
-                              settings_namespace::INVALID);
+  EXTENSION_FUNCTION_PRERUN_VALIDATE(settings_namespace_ !=
+                                     settings_namespace::INVALID);
 
   if (extension()->is_login_screen_extension() &&
       storage_area_ != StorageAreaNamespace::kManaged) {
     // Login screen extensions are not allowed to use local/sync storage for
     // security reasons (see crbug.com/978443).
-    return RespondNow(Error(base::StringPrintf(
+    *error = base::StringPrintf(
         "\"%s\" is not available for login screen extensions",
-        storage_area_string.c_str())));
+        storage_area_string.c_str());
+    return false;
   }
 
   StorageFrontend* frontend = StorageFrontend::Get(browser_context());
   if (!frontend->IsStorageEnabled(settings_namespace_)) {
-    return RespondNow(Error(
+    *error =
         base::StringPrintf("\"%s\" is not available in this instance of Chrome",
-                           storage_area_string.c_str())));
+                           storage_area_string.c_str());
+    return false;
   }
 
+  return true;
+}
+
+ExtensionFunction::ResponseAction SettingsFunction::Run() {
+  if (storage_area_ == StorageAreaNamespace::kSession) {
+    return RespondNow(RunInSession());
+  }
+
+  StorageFrontend* frontend = StorageFrontend::Get(browser_context());
+
   observer_ = GetSequenceBoundSettingsChangedCallback(
       base::SequencedTaskRunner::GetCurrentDefault(), frontend->GetObserver());
 
@@ -201,6 +218,21 @@ ExtensionFunction::ResponseAction SettingsFunction::Run() {
   return RespondLater();
 }
 
+ExtensionFunction::ResponseValue SettingsFunction::RunWithStorage(
+    ValueStore* storage) {
+  // TODO(crbug.com/40963428): Remove this when RunWithStorage has been removed
+  // from all functions.
+  NOTREACHED();
+  return BadMessage();
+}
+
+ExtensionFunction::ResponseValue SettingsFunction::RunInSession() {
+  // TODO(crbug.com/40963428): Remove this when RunWithStorage has been removed
+  // from all functions.
+  NOTREACHED();
+  return BadMessage();
+}
+
 void SettingsFunction::AsyncRunWithStorage(ValueStore* storage) {
   ResponseValue response = RunWithStorage(storage);
   content::GetUIThreadTaskRunner({})->PostTask(
@@ -345,70 +377,53 @@ ExtensionFunction::ResponseValue StorageStorageAreaGetFunction::RunInSession() {
   return WithArguments(std::move(value_dict));
 }
 
-ExtensionFunction::ResponseValue
-StorageStorageAreaGetBytesInUseFunction::RunWithStorage(ValueStore* storage) {
-  TRACE_EVENT1("browser",
-               "StorageStorageAreaGetBytesInUseFunction::RunWithStorage",
-               "extension_id", extension_id());
-
-  if (args().empty())
-    return BadMessage();
-  const base::Value& input = args()[0];
-
-  size_t bytes_in_use = 0;
-
-  switch (input.type()) {
-    case base::Value::Type::NONE:
-      bytes_in_use = storage->GetBytesInUse();
-      break;
-
-    case base::Value::Type::STRING:
-      bytes_in_use = storage->GetBytesInUse(input.GetString());
-      break;
-
-    case base::Value::Type::LIST:
-      bytes_in_use = storage->GetBytesInUse(GetKeysFromList(input.GetList()));
-      break;
-
-    default:
-      return BadMessage();
+ExtensionFunction::ResponseAction
+StorageStorageAreaGetBytesInUseFunction::Run() {
+  if (args().empty()) {
+    return RespondNow(BadMessage());
   }
 
-  return WithArguments(static_cast<double>(bytes_in_use));
-}
-
-ExtensionFunction::ResponseValue
-StorageStorageAreaGetBytesInUseFunction::RunInSession() {
-  if (args().empty())
-    return BadMessage();
   const base::Value& input = args()[0];
-
-  size_t bytes_in_use = 0;
-  SessionStorageManager* session_manager =
-      SessionStorageManager::GetForBrowserContext(browser_context());
+  std::optional<std::vector<std::string>> keys;
 
   switch (input.type()) {
     case base::Value::Type::NONE:
-      bytes_in_use = session_manager->GetTotalBytesInUse(extension_id());
+      keys = std::nullopt;
       break;
 
     case base::Value::Type::STRING:
-      bytes_in_use = session_manager->GetBytesInUse(
-          extension_id(), std::vector<std::string>(1, input.GetString()));
+      keys = std::optional(std::vector<std::string>(1, input.GetString()));
       break;
 
     case base::Value::Type::LIST:
-      bytes_in_use = session_manager->GetBytesInUse(
-          extension_id(), GetKeysFromList(input.GetList()));
+      keys = std::optional(GetKeysFromList(input.GetList()));
       break;
 
     default:
-      return BadMessage();
+      return RespondNow(BadMessage());
+  }
+
+  StorageFrontend* frontend = StorageFrontend::Get(browser_context());
+  frontend->GetBytesInUse(
+      extension(), storage_area(), keys,
+      base::BindOnce(&StorageStorageAreaGetBytesInUseFunction::
+                         OnGetBytesInUseOperationFinished,
+                     this));
+
+  return RespondLater();
+}
+
+void StorageStorageAreaGetBytesInUseFunction::OnGetBytesInUseOperationFinished(
+    size_t bytes_in_use) {
+  // Since the storage access happens asynchronously, the browser context can
+  // be torn down in the interim. If this happens, early-out.
+  if (!browser_context()) {
+    return;
   }
 
   // Checked cast should not overflow since a double can represent up to 2*53
   // bytes before a loss of precision.
-  return WithArguments(base::checked_cast<double>(bytes_in_use));
+  Respond(WithArguments(base::checked_cast<double>(bytes_in_use)));
 }
 
 ExtensionFunction::ResponseValue StorageStorageAreaSetFunction::RunWithStorage(

extensions/browser/api/storage/storage_api.h

@@ -11,6 +11,7 @@
 #include "extensions/browser/api/storage/settings_namespace.h"
 #include "extensions/browser/api/storage/settings_observer.h"
 #include "extensions/browser/api/storage/storage_area_namespace.h"
+#include "extensions/browser/api/storage/storage_frontend.h"
 #include "extensions/browser/extension_function.h"
 
 namespace extensions {
@@ -23,15 +24,16 @@ class SettingsFunction : public ExtensionFunction {
 
   // ExtensionFunction:
   bool ShouldSkipQuotaLimiting() const override;
+  bool PreRunValidation(std::string* error) override;
   ResponseAction Run() override;
 
   // Extension settings function implementations should do their work here.
   // The StorageFrontend makes sure this is posted to the appropriate thread.
-  virtual ResponseValue RunWithStorage(value_store::ValueStore* storage) = 0;
+  virtual ResponseValue RunWithStorage(value_store::ValueStore* storage);
 
   // Extension settings function implementations in `session` namespace should
   // do their work here.
-  virtual ResponseValue RunInSession() = 0;
+  virtual ResponseValue RunInSession();
 
   // Convert the |result| of a read function to the appropriate response value.
   // - If the |result| succeeded this will return a response object argument.
@@ -51,6 +53,8 @@ class SettingsFunction : public ExtensionFunction {
   // Returns whether the caller's context has access to the storage or not.
   bool IsAccessToStorageAllowed();
 
+  StorageAreaNamespace storage_area() const { return storage_area_; }
+
  private:
   // Called via PostTask from Run. Calls RunWithStorage and then
   // SendResponse with its success value.
@@ -136,8 +140,10 @@ class StorageStorageAreaGetBytesInUseFunction : public SettingsFunction {
   ~StorageStorageAreaGetBytesInUseFunction() override {}
 
   // SettingsFunction:
-  ResponseValue RunWithStorage(value_store::ValueStore* storage) override;
-  ResponseValue RunInSession() override;
+  ResponseAction Run() override;
+
+  // Called after retrieving bytes from storage.
+  void OnGetBytesInUseOperationFinished(size_t);
 };
 
 class StorageStorageAreaSetAccessLevelFunction : public SettingsFunction {

extensions/browser/api/storage/storage_api_unittest.cc

@@ -4,7 +4,7 @@
 
 #include "extensions/browser/api/storage/storage_api.h"
 
-#include "stdint.h"
+#include <stdint.h>
 
 #include <limits>
 #include <memory>
@@ -17,6 +17,7 @@
 #include "base/logging.h"
 #include "base/memory/ref_counted.h"
 #include "base/strings/stringprintf.h"
+#include "base/test/bind.h"
 #include "components/crx_file/id_util.h"
 #include "components/value_store/leveldb_value_store.h"
 #include "components/value_store/value_store.h"
@@ -27,9 +28,12 @@
 #include "extensions/browser/api/storage/settings_storage_quota_enforcer.h"
 #include "extensions/browser/api/storage/settings_test_util.h"
 #include "extensions/browser/api/storage/storage_frontend.h"
+#include "extensions/browser/api/storage/value_store_cache.h"
+#include "extensions/browser/api_test_utils.h"
 #include "extensions/browser/api_unittest.h"
 #include "extensions/browser/event_router.h"
 #include "extensions/browser/event_router_factory.h"
+#include "extensions/browser/extension_function.h"
 #include "extensions/browser/test_event_router_observer.h"
 #include "extensions/browser/test_extensions_browser_client.h"
 #include "extensions/common/api/storage.h"
@@ -282,6 +286,26 @@ TEST_F(StorageApiUnittest, GetBytesInUseIntOverflow) {
     size_t bytes_in_use_ = 0;
   };
 
+  // Create a fake ValueStoreCache that we can assign to a storage area in the
+  // StorageFrontend. This allows us to call StorageFrontend using an extension
+  // API and access our mock ValueStore.
+  class FakeValueStoreCache : public ValueStoreCache {
+   public:
+    explicit FakeValueStoreCache(FakeValueStore store) : store_(store) {}
+
+    // ValueStoreCache:
+    void ShutdownOnUI() override {}
+    void RunWithValueStoreForExtension(
+        FakeValueStoreCache::StorageCallback callback,
+        scoped_refptr<const Extension> extension) override {
+      std::move(callback).Run(&store_);
+    }
+    void DeleteStorageSoon(const ExtensionId& extension_id) override {}
+
+   private:
+    FakeValueStore store_;
+  };
+
   static constexpr struct TestCase {
     size_t bytes_in_use;
     double result;
@@ -293,21 +317,28 @@ TEST_F(StorageApiUnittest, GetBytesInUseIntOverflow) {
       {static_cast<size_t>(std::numeric_limits<int>::max()) + 1,
        static_cast<size_t>(std::numeric_limits<int>::max()) + 1}};
 
+  StorageFrontend* frontend = StorageFrontend::Get(browser_context());
+
   for (const auto& test_case : test_cases) {
     FakeValueStore value_store(test_case.bytes_in_use);
+    frontend->SetCacheForTesting(
+        settings_namespace::Namespace::LOCAL,
+        std::make_unique<FakeValueStoreCache>(value_store));
+
     auto function =
         base::MakeRefCounted<StorageStorageAreaGetBytesInUseFunction>();
-    // Call into the protected member function to avoid dealing with how the
-    // base class gets the StorageArea to use. Set `args` to a base::Value
-    // "none" so we get the total bytes in use.
-    base::Value::List args;
-    args.Append(base::Value());
-    function->SetArgs(std::move(args));
-    function->RunWithStorage(&value_store);
-    const base::Value::List* results = function->GetResultListForTest();
-    ASSERT_TRUE(results);
-    ASSERT_EQ(1u, results->size());
-    EXPECT_EQ(test_case.result, (*results)[0]);
+
+    function->set_extension(extension());
+
+    std::optional<base::Value> result =
+        api_test_utils::RunFunctionAndReturnSingleResult(
+            function.get(),
+            base::Value::List().Append("local").Append(base::Value()),
+            browser_context());
+    ASSERT_TRUE(result);
+    ASSERT_TRUE(result->is_double());
+    EXPECT_EQ(test_case.result, result->GetDouble());
+    frontend->DisableStorageForTesting(settings_namespace::Namespace::LOCAL);
   }
 }
 

extensions/browser/api/storage/storage_frontend.cc

@@ -35,6 +35,7 @@
 
 using content::BrowserContext;
 using content::BrowserThread;
+using value_store::ValueStore;
 
 namespace extensions {
 
@@ -60,6 +61,16 @@ events::HistogramValue StorageAreaToEventHistogram(
   }
 }
 
+void GetBytesInUseWithValueStore(std::optional<std::vector<std::string>> keys,
+                                 base::OnceCallback<void(size_t)> callback,
+                                 ValueStore* store) {
+  size_t size = keys.has_value() ? store->GetBytesInUse(keys.value())
+                                 : store->GetBytesInUse();
+
+  content::GetUIThreadTaskRunner({})->PostTask(
+      FROM_HERE, base::BindOnce(std::move(callback), size));
+}
+
 }  // namespace
 
 // static
@@ -110,6 +121,39 @@ StorageFrontend::~StorageFrontend() {
   }
 }
 
+void StorageFrontend::GetBytesInUse(
+    scoped_refptr<const Extension> extension,
+    StorageAreaNamespace storage_area,
+    std::optional<std::vector<std::string>> keys,
+    base::OnceCallback<void(size_t)> callback) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+
+  if (storage_area == StorageAreaNamespace::kSession) {
+    SessionStorageManager* storage_manager =
+        SessionStorageManager::GetForBrowserContext(browser_context_);
+
+    size_t bytes_in_use =
+        keys.has_value()
+            ? storage_manager->GetBytesInUse(extension->id(), keys.value())
+            : storage_manager->GetTotalBytesInUse(extension->id());
+
+    // Using a task here is important since we want to consistently fire the
+    // callback asynchronously.
+    content::GetUIThreadTaskRunner({})->PostTask(
+        FROM_HERE, base::BindOnce(std::move(callback), bytes_in_use));
+    return;
+  }
+
+  extensions::settings_namespace::Namespace settings_namespace =
+      extensions::StorageAreaToSettingsNamespace(storage_area);
+
+  CHECK(StorageFrontend::IsStorageEnabled(settings_namespace));
+
+  RunWithStorage(extension, settings_namespace,
+                 base::BindOnce(&GetBytesInUseWithValueStore, std::move(keys),
+                                std::move(callback)));
+}
+
 ValueStoreCache* StorageFrontend::GetValueStoreCache(
     settings_namespace::Namespace settings_namespace) const {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
@@ -161,6 +205,13 @@ SettingsChangedCallback StorageFrontend::GetObserver() {
                              weak_factory_.GetWeakPtr());
 }
 
+void StorageFrontend::SetCacheForTesting(
+    settings_namespace::Namespace settings_namespace,
+    std::unique_ptr<ValueStoreCache> cache) {
+  DisableStorageForTesting(settings_namespace);  // IN-TEST
+  caches_[settings_namespace] = cache.release();
+}
+
 void StorageFrontend::DisableStorageForTesting(
     settings_namespace::Namespace settings_namespace) {
   auto it = caches_.find(settings_namespace);

extensions/browser/api/storage/storage_frontend.h

@@ -13,6 +13,8 @@
 #include "base/memory/scoped_refptr.h"
 #include "base/memory/weak_ptr.h"
 #include "base/values.h"
+#include "components/value_store/value_store.h"
+#include "extensions/browser/api/storage/session_storage_manager.h"
 #include "extensions/browser/api/storage/settings_namespace.h"
 #include "extensions/browser/api/storage/settings_observer.h"
 #include "extensions/browser/api/storage/value_store_cache.h"
@@ -65,9 +67,21 @@ class StorageFrontend : public BrowserContextKeyedAPI {
   void DeleteStorageSoon(const ExtensionId& extension_id,
                          base::OnceClosure done_callback);
 
+  // For a given `extension` and `storage_area`, determines the number of bytes
+  // in use and fires `callback` with the result. If `keys` is specified, the
+  // result is based only on keys contained within the vector. Otherwise, all
+  // keys are included.
+  void GetBytesInUse(scoped_refptr<const Extension> extension,
+                     StorageAreaNamespace storage_area,
+                     std::optional<std::vector<std::string>> keys,
+                     base::OnceCallback<void(size_t)> callback);
+
   // Gets the Settings change callback.
   SettingsChangedCallback GetObserver();
 
+  void SetCacheForTesting(settings_namespace::Namespace settings_namespace,
+                          std::unique_ptr<ValueStoreCache> cache);
+
   void DisableStorageForTesting(
       settings_namespace::Namespace settings_namespace);