diff --git a/Source/WebKit/UIProcess/Extensions/Cocoa/API/WebExtensionContextAPIStorageCocoa.mm b/Source/WebKit/UIProcess/Extensions/Cocoa/API/WebExtensionContextAPIStorageCocoa.mm index 080963434926e..b123d4bd0efeb 100644 --- a/Source/WebKit/UIProcess/Extensions/Cocoa/API/WebExtensionContextAPIStorageCocoa.mm +++ b/Source/WebKit/UIProcess/Extensions/Cocoa/API/WebExtensionContextAPIStorageCocoa.mm @@ -55,7 +55,7 @@ { auto callingAPIName = makeString("browser.storage."_s, toAPIString(dataType), ".get()"_s); - auto storage = storageForType(dataType); + auto *storage = storageForType(dataType); [storage getValuesForKeys:createNSArray(keys).get() completionHandler:makeBlockPtr([callingAPIName, completionHandler = WTFMove(completionHandler)](NSDictionary *values, NSString *errorMessage) mutable { if (errorMessage) completionHandler(toWebExtensionError(callingAPIName, nil, errorMessage)); @@ -64,11 +64,24 @@ }).get()]; } +void WebExtensionContext::storageGetKeys(WebPageProxyIdentifier webPageProxyIdentifier, WebExtensionDataType dataType, CompletionHandler, WebExtensionError>&&)>&& completionHandler) +{ + auto callingAPIName = makeString("browser.storage."_s, toAPIString(dataType), ".getKeys()"_s); + + auto *storage = storageForType(dataType); + [storage getAllKeys:makeBlockPtr([callingAPIName, completionHandler = WTFMove(completionHandler)](NSArray *keys, NSString *errorMessage) mutable { + if (errorMessage) + completionHandler(toWebExtensionError(callingAPIName, nil, errorMessage)); + else + completionHandler(makeVector(keys)); + }).get()]; +} + void WebExtensionContext::storageGetBytesInUse(WebPageProxyIdentifier webPageProxyIdentifier, WebExtensionDataType dataType, const Vector& keys, CompletionHandler&&)>&& completionHandler) { auto callingAPIName = makeString("browser.storage."_s, toAPIString(dataType), ".getBytesInUse()"_s); - auto storage = storageForType(dataType); + auto *storage = storageForType(dataType); [storage getStorageSizeForKeys:createNSArray(keys).get() completionHandler:makeBlockPtr([callingAPIName, completionHandler = WTFMove(completionHandler)](size_t size, NSString *errorMessage) mutable { if (errorMessage) completionHandler(toWebExtensionError(callingAPIName, nil, errorMessage)); diff --git a/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.h b/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.h index 294990c2e3fb1..64fa06f10b108 100644 --- a/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.h +++ b/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.h @@ -40,6 +40,7 @@ enum class WebExtensionDataType : uint8_t; - (instancetype)initWithUniqueIdentifier:(NSString *)uniqueIdentifier storageType:(WebKit::WebExtensionDataType)storageType directory:(NSString *)directory usesInMemoryDatabase:(BOOL)useInMemoryDatabase; +- (void)getAllKeys:(void (^)(NSArray *keys, NSString * _Nullable errorMessage))completionHandler; - (void)getValuesForKeys:(NSArray *)keys completionHandler:(void (^)(NSDictionary *results, NSString * _Nullable errorMessage))completionHandler; - (void)getStorageSizeForKeys:(NSArray *)keys completionHandler:(void (^)(size_t storageSize, NSString * _Nullable errorMessage))completionHandler; - (void)getStorageSizeForAllKeysIncludingKeyedData:(NSDictionary *)additionalKeyedData withCompletionHandler:(void (^)(size_t storageSize, NSUInteger numberOfKeysIncludingAdditionalKeyedData, NSDictionary *existingKeysAndValues, NSString * _Nullable errorMessage))completionHandler; diff --git a/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.mm b/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.mm index 66a6a597c8bfd..e05766352eca8 100644 --- a/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.mm +++ b/Source/WebKit/UIProcess/Extensions/Cocoa/_WKWebExtensionStorageSQLiteStore.mm @@ -71,14 +71,34 @@ - (instancetype)initWithUniqueIdentifier:(NSString *)uniqueIdentifier storageTyp return self; } +- (void)getAllKeys:(void (^)(NSArray *keys, NSString *errorMessage))completionHandler +{ + auto weakSelf = WeakObjCPtr { self }; + dispatch_async(_databaseQueue, ^{ + auto *strongSelf = weakSelf.get().get(); + if (!strongSelf) { + RELEASE_LOG_ERROR(Extensions, "Failed to retrieve all keys for extension %{private}@.", self->_uniqueIdentifier); + completionHandler(nil, @"Failed to retrieve all keys"); + return; + } + + NSString *errorMessage; + auto *keysArray = [self _getAllKeysReturningErrorMessage:&errorMessage].allObjects; + + dispatch_async(dispatch_get_main_queue(), ^{ + completionHandler(keysArray, errorMessage); + }); + }); +} + - (void)getValuesForKeys:(NSArray *)keys completionHandler:(void (^)(NSDictionary *results, NSString *errorMessage))completionHandler { - auto weakSelf = WeakObjCPtr<_WKWebExtensionStorageSQLiteStore> { self }; + auto weakSelf = WeakObjCPtr { self }; dispatch_async(_databaseQueue, ^{ - auto strongSelf = weakSelf.get(); + auto *strongSelf = weakSelf.get().get(); if (!strongSelf) { - RELEASE_LOG_ERROR(Extensions, "Failed to retrieve keys: %{private}@ for extension %{private}@.", keys, self->_uniqueIdentifier); - completionHandler(nil, [NSString stringWithFormat:@"Failed to retrieve keys %@", keys]); + RELEASE_LOG_ERROR(Extensions, "Failed to retrieve values for keys: %{private}@ for extension %{private}@.", keys, self->_uniqueIdentifier); + completionHandler(nil, [NSString stringWithFormat:@"Failed to retrieve values for keys %@", keys]); return; } @@ -93,9 +113,9 @@ - (void)getValuesForKeys:(NSArray *)keys completionHandler:(void (^) - (void)getStorageSizeForKeys:(NSArray *)keys completionHandler:(void (^)(size_t storageSize, NSString *errorMessage))completionHandler { - auto weakSelf = WeakObjCPtr<_WKWebExtensionStorageSQLiteStore> { self }; + auto weakSelf = WeakObjCPtr { self }; dispatch_async(_databaseQueue, ^{ - auto strongSelf = weakSelf.get(); + auto *strongSelf = weakSelf.get().get(); if (!strongSelf) { RELEASE_LOG_ERROR(Extensions, "Failed to calculate storage size for keys: %{private}@ for extension %{private}@.", keys, self->_uniqueIdentifier); completionHandler(0, [NSString stringWithFormat:@"Failed to caluclate storage size for keys: %@", keys]); @@ -144,9 +164,9 @@ - (void)getStorageSizeForAllKeysIncludingKeyedData:(NSDictionary { self }; + auto weakSelf = WeakObjCPtr { self }; dispatch_async(self->_databaseQueue, ^{ - auto strongSelf = weakSelf.get(); + auto *strongSelf = weakSelf.get().get(); if (!strongSelf) { RELEASE_LOG_ERROR(Extensions, "Failed to calculate storage size for extension %{private}@.", self->_uniqueIdentifier); completionHandler(0.0, 0, @{ }, @"Failed to calculate storage size"); @@ -171,9 +191,9 @@ - (void)getStorageSizeForAllKeysIncludingKeyedData:(NSDictionary *)keyedData completionHandler:(void (^)(NSArray *keysSuccessfullySet, NSString *errorMessage))completionHandler { - auto weakSelf = WeakObjCPtr<_WKWebExtensionStorageSQLiteStore> { self }; + auto weakSelf = WeakObjCPtr { self }; dispatch_async(_databaseQueue, ^{ - auto strongSelf = weakSelf.get(); + auto *strongSelf = weakSelf.get().get(); if (!strongSelf) { completionHandler(nil, [NSString stringWithFormat:@"Failed to set keys %@", keyedData.allKeys]); return; @@ -207,9 +227,9 @@ - (void)setKeyedData:(NSDictionary *)keyedData completio - (void)deleteValuesForKeys:(NSArray *)keys completionHandler:(void (^)(NSString *errorMessage))completionHandler { - auto weakSelf = WeakObjCPtr<_WKWebExtensionStorageSQLiteStore> { self }; + auto weakSelf = WeakObjCPtr { self }; dispatch_async(_databaseQueue, ^{ - auto strongSelf = weakSelf.get(); + auto *strongSelf = weakSelf.get().get(); if (!strongSelf) { completionHandler([NSString stringWithFormat:@"Failed to delete keys %@", keys]); return; diff --git a/Source/WebKit/UIProcess/Extensions/WebExtensionContext.h b/Source/WebKit/UIProcess/Extensions/WebExtensionContext.h index 828490c95dfbc..015ebb3d86244 100644 --- a/Source/WebKit/UIProcess/Extensions/WebExtensionContext.h +++ b/Source/WebKit/UIProcess/Extensions/WebExtensionContext.h @@ -857,6 +857,7 @@ class WebExtensionContext : public API::ObjectImpl& keys, CompletionHandler&&)>&&); + void storageGetKeys(WebPageProxyIdentifier, WebExtensionDataType, CompletionHandler, WebExtensionError>&&)>&&); void storageGetBytesInUse(WebPageProxyIdentifier, WebExtensionDataType, const Vector& keys, CompletionHandler&&)>&&); void storageSet(WebPageProxyIdentifier, WebExtensionDataType, const String& dataJSON, CompletionHandler&&)>&&); void storageRemove(WebPageProxyIdentifier, WebExtensionDataType, const Vector& keys, CompletionHandler&&)>&&); diff --git a/Source/WebKit/UIProcess/Extensions/WebExtensionContext.messages.in b/Source/WebKit/UIProcess/Extensions/WebExtensionContext.messages.in index ee1081cfac3d0..a95444bda6054 100644 --- a/Source/WebKit/UIProcess/Extensions/WebExtensionContext.messages.in +++ b/Source/WebKit/UIProcess/Extensions/WebExtensionContext.messages.in @@ -133,6 +133,7 @@ messages -> WebExtensionContext { // Storage APIs [EnabledIf='isStorageMessageAllowed()'] StorageGet(WebKit::WebPageProxyIdentifier webPageProxyIdentifier, WebKit::WebExtensionDataType dataType, Vector keys) -> (Expected result) + [EnabledIf='isStorageMessageAllowed()'] StorageGetKeys(WebKit::WebPageProxyIdentifier webPageProxyIdentifier, WebKit::WebExtensionDataType dataType) -> (Expected, WebKit::WebExtensionError> result) [EnabledIf='isStorageMessageAllowed()'] StorageGetBytesInUse(WebKit::WebPageProxyIdentifier webPageProxyIdentifier, WebKit::WebExtensionDataType dataType, Vector keys) -> (Expected result) [EnabledIf='isStorageMessageAllowed()'] StorageSet(WebKit::WebPageProxyIdentifier webPageProxyIdentifier, WebKit::WebExtensionDataType dataType, String dataJSON) -> (Expected result) [EnabledIf='isStorageMessageAllowed()'] StorageRemove(WebKit::WebPageProxyIdentifier webPageProxyIdentifier, WebKit::WebExtensionDataType dataType, Vector keys) -> (Expected result) diff --git a/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIStorageAreaCocoa.mm b/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIStorageAreaCocoa.mm index a0a40370e8302..2b890ea510730 100644 --- a/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIStorageAreaCocoa.mm +++ b/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIStorageAreaCocoa.mm @@ -116,6 +116,18 @@ }, extensionContext().identifier()); } +void WebExtensionAPIStorageArea::getKeys(WebPage& page, Ref&& callback, NSString **outExceptionString) +{ + WebProcess::singleton().sendWithAsyncReply(Messages::WebExtensionContext::StorageGetKeys(page.webPageProxyIdentifier(), m_type), [protectedThis = Ref { *this }, callback = WTFMove(callback)](Expected, WebExtensionError>&& result) { + if (!result) { + callback->reportError(result.error()); + return; + } + + callback->call(createNSArray(result.value()).get()); + }, extensionContext().identifier()); +} + void WebExtensionAPIStorageArea::getBytesInUse(WebPage& page, id keys, Ref&& callback, NSString **outExceptionString) { // Documentation: https://developer.mozilla.org/docs/Mozilla/Add-ons/WebExtensions/API/storage/StorageArea/getBytesInUse diff --git a/Source/WebKit/WebProcess/Extensions/API/WebExtensionAPIStorageArea.h b/Source/WebKit/WebProcess/Extensions/API/WebExtensionAPIStorageArea.h index 6c1777160884a..967bbe7ae7316 100644 --- a/Source/WebKit/WebProcess/Extensions/API/WebExtensionAPIStorageArea.h +++ b/Source/WebKit/WebProcess/Extensions/API/WebExtensionAPIStorageArea.h @@ -43,6 +43,7 @@ class WebExtensionAPIStorageArea : public WebExtensionAPIObject, public JSWebExt bool isPropertyAllowed(const ASCIILiteral& propertyName, WebPage*); void get(WebPage&, id items, Ref&&, NSString **outExceptionString); + void getKeys(WebPage&, Ref&&, NSString **outExceptionString); void getBytesInUse(WebPage&, id keys, Ref&&, NSString **outExceptionString); void set(WebPage&, NSDictionary *items, Ref&&, NSString **outExceptionString); void remove(WebPage&, id keys, Ref&&, NSString **outExceptionString); diff --git a/Source/WebKit/WebProcess/Extensions/Interfaces/WebExtensionAPIStorageArea.idl b/Source/WebKit/WebProcess/Extensions/Interfaces/WebExtensionAPIStorageArea.idl index 3f7ce27783e5f..615988ce6fa43 100644 --- a/Source/WebKit/WebProcess/Extensions/Interfaces/WebExtensionAPIStorageArea.idl +++ b/Source/WebKit/WebProcess/Extensions/Interfaces/WebExtensionAPIStorageArea.idl @@ -30,6 +30,7 @@ ] interface WebExtensionAPIStorageArea { [RaisesException] void get([Optional, NSObject=StopAtTopLevel, DOMString] any items, [Optional, CallbackHandler] function callback); + [RaisesException] void getKeys([Optional, CallbackHandler] function callback); [RaisesException] void getBytesInUse([Optional, NSObject, DOMString] any keys, [Optional, CallbackHandler] function callback); [RaisesException] void set([NSDictionary=StopAtTopLevel] any items, [Optional, CallbackHandler] function callback); [RaisesException] void remove([NSObject] any keys, [Optional, CallbackHandler] function callback); diff --git a/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIStorage.mm b/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIStorage.mm index fd14d63177119..e093611e5bf64 100644 --- a/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIStorage.mm +++ b/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIStorage.mm @@ -59,6 +59,8 @@ auto *backgroundScript = Util::constructScript(@[ @"browser.test.assertThrows(() => browser?.storage?.local?.get(Date.now()), /'items' value is invalid, because an object or a string or an array of strings or null is expected, but a number was provided/i)", + @"browser.test.assertThrows(() => browser?.storage?.local?.getKeys('invalid'), /'callback' value is invalid, because a function is expected/i)", + @"browser.test.assertThrows(() => browser?.storage?.local?.getBytesInUse({}), /'keys' value is invalid, because a string or an array of strings or null is expected, but an object was provided/i)", @"browser.test.assertThrows(() => browser?.storage?.local?.getBytesInUse([1]), /'keys' value is invalid, because a string or an array of strings or null is expected, but an array of other values was provided/i)", @@ -311,6 +313,36 @@ Util::loadAndRunExtension(storageManifest, @{ @"background.js": backgroundScript }); } +TEST(WKWebExtensionAPIStorage, GetKeys) +{ + auto *backgroundScript = Util::constructScript(@[ + @"const data = { 'string': 'string', 'number': 1, 'boolean': true, 'dictionary': {'key': 'value'}, 'array': [1, true, 'string'], 'null': null }", + @"await browser?.storage?.local?.set(data)", + + @"var keys = await browser?.storage?.local?.getKeys()", + @"browser.test.assertEq(keys.length, 6, 'Should have 6 keys')", + @"browser.test.assertTrue(keys.includes('string'), 'Should include string key')", + @"browser.test.assertTrue(keys.includes('number'), 'Should include number key')", + @"browser.test.assertTrue(keys.includes('boolean'), 'Should include boolean key')", + @"browser.test.assertTrue(keys.includes('dictionary'), 'Should include dictionary key')", + @"browser.test.assertTrue(keys.includes('array'), 'Should include array key')", + @"browser.test.assertTrue(keys.includes('null'), 'Should include null key')", + + @"await browser?.storage?.local?.remove('number')", + @"keys = await browser?.storage?.local?.getKeys()", + @"browser.test.assertEq(keys.length, 5, 'Should have 5 keys after removal')", + @"browser.test.assertFalse(keys.includes('number'), 'Should not include removed number key')", + + @"await browser?.storage?.local?.clear()", + @"keys = await browser?.storage?.local?.getKeys()", + @"browser.test.assertEq(keys.length, 0, 'Should have no keys after clear')", + + @"browser.test.notifyPass()", + ]); + + Util::loadAndRunExtension(storageManifest, @{ @"background.js": backgroundScript }); +} + TEST(WKWebExtensionAPIStorage, GetBytesInUse) { auto *backgroundScript = Util::constructScript(@[