[Extensions] Use GetKeys() ValueStore method in extension API.

Use the GetKeys() method implemented in [1] for the implementation of
the `chrome.storage` API. Also changes the implementation of the session
storage area to support reading keys without reading values.

[1] https://chromium-review.googlesource.com/c/chromium/src/+/6197397

Bug: 380739760
Change-Id: I463661e5deff8e595a437174eb80ff2e74b1fbb1
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6198270
Commit-Queue: Oliver Dunk <oliverdunk@chromium.org>
Reviewed-by: Emilia Paz <emiliapaz@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1411823}

extensions/browser/api/storage/session_storage_manager.cc

@@ -123,6 +123,15 @@ size_t SessionStorageManager::ExtensionStorage::CalculateUsage(
   return updated_used_total;
 }
 
+std::vector<std::string> SessionStorageManager::ExtensionStorage::GetKeys()
+    const {
+  std::vector<std::string> keys;
+  for (auto& value : values_) {
+    keys.push_back(value.first);
+  }
+  return keys;
+}
+
 std::map<std::string, const base::Value*>
 SessionStorageManager::ExtensionStorage::Get(
     const std::vector<std::string>& keys) const {
@@ -251,6 +260,16 @@ BrowserContextKeyedServiceFactory* SessionStorageManager::GetFactory() {
   return g_factory.get();
 }
 
+std::vector<std::string> SessionStorageManager::GetKeys(
+    const ExtensionId& extension_id) const {
+  auto storage_it = extensions_storage_.find(extension_id);
+  if (storage_it == extensions_storage_.end()) {
+    return std::vector<std::string>();
+  }
+
+  return storage_it->second->GetKeys();
+}
+
 const base::Value* SessionStorageManager::Get(const ExtensionId& extension_id,
                                               const std::string& key) const {
   std::map<std::string, const base::Value*> values =

extensions/browser/api/storage/session_storage_manager.h

@@ -60,6 +60,10 @@ class SessionStorageManager : public KeyedService,
   // Retrieves the factory instance for the SessionStorageManager.
   static BrowserContextKeyedServiceFactory* GetFactory();
 
+  // Returns a vector with all keys found in storage for the given
+  // `extension_id`.
+  std::vector<std::string> GetKeys(const ExtensionId& extension_id) const;
+
   // Returns the value for the given `extension_id` and `key`, or null if none
   // exists.
   const base::Value* Get(const ExtensionId& extension_id,
@@ -133,6 +137,9 @@ class SessionStorageManager : public KeyedService,
     explicit ExtensionStorage(size_t quota_bytes);
     ~ExtensionStorage();
 
+    // Returns a vector with keys found in storage.
+    std::vector<std::string> GetKeys() const;
+
     // Returns a map with keys and values found in storage.
     std::map<std::string, const base::Value*> Get(
         const std::vector<std::string>& keys) const;

extensions/browser/api/storage/storage_frontend.cc

@@ -60,6 +60,15 @@ events::HistogramValue StorageAreaToEventHistogram(
   }
 }
 
+void GetKeysWithValueStore(
+    base::OnceCallback<void(ValueStore::ReadResult)> callback,
+    ValueStore* store) {
+  ValueStore::ReadResult result = store->GetKeys();
+
+  content::GetUIThreadTaskRunner({})->PostTask(
+      FROM_HERE, base::BindOnce(std::move(callback), std::move(result)));
+}
+
 void GetWithValueStore(
     std::optional<std::vector<std::string>> keys,
     base::OnceCallback<void(ValueStore::ReadResult)> callback,
@@ -213,14 +222,18 @@ void StorageFrontend::OnReadFinished(
 
 void StorageFrontend::OnReadKeysFinished(
     base::OnceCallback<void(GetKeysResult)> callback,
-    GetResult get_result) {
+    ValueStore::ReadResult result) {
+  bool success = result.status().ok();
+
   GetKeysResult get_keys_result;
 
-  get_keys_result.status = get_result.status;
-  get_keys_result.data =
-      get_result.status.success
-          ? std::optional(KeysFromDict(std::move(*get_result.data)))
-          : std::nullopt;
+  get_keys_result.status.success = success;
+  get_keys_result.status.error =
+      success ? std::nullopt : std::optional(result.status().message);
+
+  if (success) {
+    get_keys_result.data = KeysFromDict(result.PassSettings());
+  }
 
   std::move(callback).Run(std::move(get_keys_result));
 }
@@ -290,9 +303,40 @@ void StorageFrontend::GetKeys(
     scoped_refptr<const Extension> extension,
     StorageAreaNamespace storage_area,
     base::OnceCallback<void(GetKeysResult)> callback) {
-  GetValues(extension, storage_area, std::nullopt,
-            base::BindOnce(&StorageFrontend::OnReadKeysFinished,
-                           weak_factory_.GetWeakPtr(), std::move(callback)));
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+
+  if (storage_area == StorageAreaNamespace::kSession) {
+    SessionStorageManager* storage_manager =
+        SessionStorageManager::GetForBrowserContext(browser_context_);
+
+    std::vector<std::string> keys = storage_manager->GetKeys(extension->id());
+
+    base::Value::List list = base::Value::List::with_capacity(keys.size());
+    for (std::string key : keys) {
+      list.Append(key);
+    }
+
+    GetKeysResult get_keys_result;
+    get_keys_result.data = std::move(list);
+
+    // 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), std::move(get_keys_result)));
+    return;
+  }
+
+  settings_namespace::Namespace settings_namespace =
+      StorageAreaToSettingsNamespace(storage_area);
+
+  CHECK(StorageFrontend::IsStorageEnabled(settings_namespace));
+
+  base::OnceCallback<void(ValueStore::ReadResult)> test =
+      base::BindOnce(&StorageFrontend::OnReadKeysFinished,
+                     weak_factory_.GetWeakPtr(), std::move(callback));
+  RunWithStorage(extension, settings_namespace,
+                 base::BindOnce(&GetKeysWithValueStore, std::move(test)));
 }
 
 void StorageFrontend::GetBytesInUse(

extensions/browser/api/storage/storage_frontend.h

@@ -170,9 +170,9 @@ class StorageFrontend : public BrowserContextKeyedAPI {
   void Init(scoped_refptr<value_store::ValueStoreFactory> storage_factory);
 
   // Should be called on the UI thread after a read has been performed in
-  // `storage_area`. Fires `callback` with the keys from `get_result`.
+  // `storage_area`. Fires `callback` with the keys from `result`.
   void OnReadKeysFinished(base::OnceCallback<void(GetKeysResult)> callback,
-                          GetResult get_result);
+                          value_store::ValueStore::ReadResult result);
 
   // Should be called on the UI thread after a read has been performed in
   // `storage_area`. Fires `callback` with the `result` from the read