diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 3d36ecb2b78..1f90f40d915 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -16,6 +16,7 @@ set(_coreai_runtime_sources runtime/coreai_cache_root.mm runtime/coreai_fault_injection.mm runtime/coreai_manifest.mm + runtime/coreai_source_bundle.mm runtime/coreai_validation.mm runtime/coreai_storage.mm ) @@ -44,7 +45,9 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") runtime/test/coreai_cache_key_test.mm runtime/test/coreai_bookmark_test.mm runtime/test/coreai_bookmark_contention_test.mm + runtime/test/coreai_source_test.mm runtime/test/coreai_bookmark_fixture.mm + runtime/test/coreai_source_fixture.mm ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index f4f90c49708..22e67e58524 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -56,11 +56,46 @@ The assets root must be an absolute path. The application chooses it, should reserve it for this backend, and is trusted not to rename, replace or modify its contents while the backend uses it; the backend does not defend against other processes of the same user rebinding paths. Written files and changed directories -are synced with `fsync`. Atomic publication uses `F_FULLFSYNC` on the temporary -file before rename and on the directory afterward. A failure before rename -preserves the old contents; a failure of the final flush reports an error even -though the new complete contents are already visible. Use a local APFS volume; -unsupported full-device flushes are reported as errors, not downgraded to `fsync`. +are synced with `fsync`. Bookmark publication uses `F_FULLFSYNC` on the temporary +file before rename and on the directory afterward. Source-bundle publication +fully flushes the staging directory after syncing its files and subdirectories, +then fully flushes the destination directory after rename. A failure before +rename preserves the old contents; a failure of the final flush reports an error +even though the new complete contents are already visible. Use a local APFS +volume; unsupported full-device flushes are reported as errors, not downgraded +to `fsync`. + +Inline bundles are read through `NamedDataMap` and materialized as unchanged +files under the assets root. Named data can be supplied externally; inline +packaging does not require all bytes to reside in one physical PTE. The exporter +computes each bundle digest from all delivered relative filenames and file bytes, +not just the SDK's bytecode-only `main.hash`. Each AOT architecture has an +independent digest; only the selected architecture is materialized. NDS key names +and raw asset bytes are unchanged. + +When source recovery is needed, existing bundles are checked for the expected +file set, sizes, entry types and directory structure without reading asset +contents or requesting NDS payloads. Unicode paths are compared in canonical +decomposed form; NDS keys and bytes remain unchanged. Reuse fully flushes the +source root, including after a previous publication's final flush failed. +Missing bundles are written to a staging directory and published with an +exclusive rename; each selected NDS payload is fetched once and checked for size +before writing. If another loader publishes first, its bundle is validated and +used. Incomplete, size-mismatched or otherwise malformed existing bundles fail +normal loading without replacement, since an SDK model may still be using them. +A cache-maintenance caller holding the same entry lock can explicitly replace a +damaged bundle. Replacement preflights the full descriptor-bound tree, then +revalidates entries while removing them without following links. Only +interrupted system calls are retried; other failures are returned to the caller +without automatic retries or repair. Temporary bundles are removed on ordinary +failures before publication; a caller can explicitly retry after addressing the +failure. + +These are completeness checks, not full content-integrity verification. Same-size +content changes are not detected, and a complete stored bundle does not cause its +NDS source to be reread. The export-time digest is a trusted artifact identifier; +the runtime does not rehash source or stored bytes. Core AI may reject invalid +contents, but SDK load failures do not automatically evict the extracted copy. Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers everything beneath it. Ancestors are not modified. This is backup exclusion, not diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h index 23b364b45dc..bdbb094a120 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.h +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -33,6 +33,10 @@ class CacheEntryLock final { CacheEntryPersistence persistence() const { return persistence_; } + NSString* key() const { + return key_; + } + // Missing is nil. // Missing is nil. Empty, oversized, inaccessible, or nonregular files fail // closed. runtime::Result read_bookmark() const; diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.h b/backends/apple/coreai/runtime/coreai_fault_injection.h index 84947db16c8..f8d1cbd412f 100644 --- a/backends/apple/coreai/runtime/coreai_fault_injection.h +++ b/backends/apple/coreai/runtime/coreai_fault_injection.h @@ -16,6 +16,9 @@ enum class FaultPoint { SyncDirectory, FullSync, Rename, + Remove, + CreateStagingDirectory, + OpenStagingDirectory, BeforeCreateDirectory, BeforeCacheEntryLock, }; diff --git a/backends/apple/coreai/runtime/coreai_source_bundle.h b/backends/apple/coreai/runtime/coreai_source_bundle.h new file mode 100644 index 00000000000..dabaeef5c62 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_source_bundle.h @@ -0,0 +1,33 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#import +#include +#include + +namespace executorch::backends::coreai { + +class CacheEntryLock; +class SelectedBundle; + +enum class CorruptBundlePolicy { + Reject, + Replace, +}; + +// cache_entry remains locked for this call. Reuse validates the exact tree +// structure and sizes, not file contents. +runtime::Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* named_data, + const CacheEntryLock& cache_entry, + CorruptBundlePolicy corrupt_policy = CorruptBundlePolicy::Reject); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_source_bundle.mm b/backends/apple/coreai/runtime/coreai_source_bundle.mm new file mode 100644 index 00000000000..916637e2e68 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_source_bundle.mm @@ -0,0 +1,360 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#import "coreai_source_bundle.h" + +#import "coreai_bookmarks.h" +#import "coreai_fault_injection.h" +#import "coreai_manifest.h" +#import "coreai_storage.h" + +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +enum class EntryState { + Missing, + Valid, + Corrupt, +}; + +Error path_error(int error) { + return error == ENOENT || error == ENOTDIR || error == ELOOP + ? Error::InvalidExternalData + : Error::AccessFailed; +} + +Error validate_directory( + int fd, + NSString* relative, + NSDictionary* files, + NSSet* expected_directories, + NSMutableSet* actual_files) { + auto children = storage_children(fd); + if (!children.ok()) { + return children.error(); + } + for (NSString* name in children.get()) { + @autoreleasepool { + NSString* path = [relative stringByAppendingPathComponent:name] + .decomposedStringWithCanonicalMapping; + struct stat info; + if (lstat_at(fd, name.fileSystemRepresentation, &info) != 0) { + return path_error(errno); + } + const bool directory = S_ISDIR(info.st_mode); + ET_CHECK_OR_RETURN_ERROR( + (directory && [expected_directories containsObject:path]) || + (S_ISREG(info.st_mode) && files[path] != nil), + InvalidExternalData, + "Unexpected Core AI bundle entry: %s", + path.UTF8String); + if (directory) { + FileDescriptor child( + open_directory_at(fd, name.fileSystemRepresentation)); + if (child.get() < 0) { + return path_error(errno); + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_directory( + child.get(), path, files, expected_directories, actual_files)); + } else { + ET_CHECK_OR_RETURN_ERROR( + static_cast(info.st_size) == + files[path].unsignedLongLongValue, + InvalidExternalData, + "Core AI asset size does not match manifest"); + [actual_files addObject:path]; + } + } + } + return Error::Ok; +} + +// Filesystem conversion decomposes Unicode without changing NamedDataMap keys. +NSDictionary* normalized_files( + NSDictionary* files) { + NSMutableDictionary* normalized = + [NSMutableDictionary dictionaryWithCapacity:files.count]; + for (NSString* file in files) { + normalized[file.decomposedStringWithCanonicalMapping] = files[file]; + } + return normalized; +} + +Error validate_tree( + int fd, + NSDictionary* selected_files) { + NSDictionary* files = normalized_files(selected_files); + NSMutableSet* expected_directories = [NSMutableSet set]; + for (NSString* file in files) { + NSString* directory = file.stringByDeletingLastPathComponent; + while (directory.length > 0) { + [expected_directories addObject:directory]; + directory = directory.stringByDeletingLastPathComponent; + } + } + NSMutableSet* actual_files = [NSMutableSet set]; + ET_CHECK_OK_OR_RETURN_ERROR( + validate_directory(fd, @"", files, expected_directories, actual_files)); + ET_CHECK_OR_RETURN_ERROR( + [actual_files isEqualToSet:[NSSet setWithArray:files.allKeys]], + InvalidExternalData, + "Core AI bundle file set does not match manifest"); + return Error::Ok; +} + +Result inspect_source_bundle( + int root_fd, + NSString* key, + NSDictionary* files) { + FileDescriptor entry( + open_directory_at(root_fd, key.fileSystemRepresentation)); + if (entry.get() < 0) { + if (errno == ENOENT) { + return EntryState::Missing; + } + ET_CHECK_OR_RETURN_ERROR( + errno == ENOTDIR || errno == ELOOP, + AccessFailed, + "Cannot inspect Core AI stored bundle"); + return EntryState::Corrupt; + } + const Error error = validate_tree(entry.get(), files); + if (error == Error::Ok) { + return EntryState::Valid; + } + if (error == Error::InvalidExternalData) { + return EntryState::Corrupt; + } + return error; +} + +Error fully_sync_directory(int fd) { + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(fd)); + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { + const int fault = fault_point(FaultPoint::FullSync); + if (fault != 0) { + errno = fault; + return -1; + } + return fcntl(fd, F_FULLFSYNC); + }) == 0, + AccessFailed, + "Cannot fully synchronize Core AI source directory"); + return Error::Ok; +} + +struct StagingDirectory { + int parent_fd; + NSString* name; + bool published = false; + + StagingDirectory(int parent_fd, NSString* name) + : parent_fd(parent_fd), name(name) {} + StagingDirectory(const StagingDirectory&) = delete; + StagingDirectory& operator=(const StagingDirectory&) = delete; + StagingDirectory(StagingDirectory&& other) + : parent_fd(other.parent_fd), + name(other.name), + published(other.published) { + other.published = true; + } + + ~StagingDirectory() { + if (published) { + return; + } + const Error error = + remove_storage_directory(parent_fd, name, StorageRemovalPolicy::Strict); + if (error != Error::Ok) { + ET_LOG(Error, "Cannot remove Core AI staging directory"); + } + } +}; + +Result create_staging_directory(int root_fd) { + NSString* name = + [@".staging-" stringByAppendingString:NSUUID.UUID.UUIDString]; + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { + const int fault = fault_point(FaultPoint::CreateStagingDirectory); + if (fault != 0) { + errno = fault; + return -1; + } + return mkdirat(root_fd, name.fileSystemRepresentation, 0777); + }) == 0, + AccessFailed, + "Cannot create Core AI staging directory"); + return StagingDirectory{root_fd, [name copy]}; +} + +Result open_staging_directory(int root_fd, NSString* name) { + const int fd = retry_eintr([&] { + const int fault = fault_point(FaultPoint::OpenStagingDirectory); + if (fault != 0) { + errno = fault; + return -1; + } + return open_directory_at(root_fd, name.fileSystemRepresentation); + }); + ET_CHECK_OR_RETURN_ERROR( + fd >= 0, AccessFailed, "Cannot open Core AI staging directory"); + return fd; +} + +Error write_staging_bundle( + int staging_fd, + const SelectedBundle& bundle, + const runtime::NamedDataMap& named_data) { + for (NSString* file in bundle.files) { + @autoreleasepool { + NSString* data_key = + [NSString stringWithFormat:@"coreai/%@/%@", bundle.model_hash, file]; + auto buffer = named_data.get_data(data_key.UTF8String); + if (!buffer.ok()) { + ET_CHECK_OR_RETURN_ERROR( + buffer.error() != Error::NotFound, + InvalidExternalData, + "Missing selected Core AI named data: %s", + data_key.UTF8String); + return buffer.error(); + } + ET_CHECK_OR_RETURN_ERROR( + buffer->size() == bundle.files[file].unsignedLongLongValue && + (buffer->size() == 0 || buffer->data() != nullptr), + InvalidExternalData, + "Invalid Core AI named data buffer"); + ET_CHECK_OK_OR_RETURN_ERROR( + write_storage_file(staging_fd, file, buffer->data(), buffer->size())); + } + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_tree(staging_fd, bundle.files)); + return fully_sync_directory(staging_fd); +} + +Result publish_staging_directory( + int root_fd, + StagingDirectory& staging, + NSString* key) { + const int result = retry_eintr([&] { + const int fault = fault_point(FaultPoint::Rename); + if (fault != 0) { + errno = fault; + return -1; + } + return renameatx_np( + root_fd, + staging.name.fileSystemRepresentation, + root_fd, + key.fileSystemRepresentation, + RENAME_EXCL); + }); + if (result == 0) { + staging.published = true; + ET_CHECK_OK_OR_RETURN_ERROR(fully_sync_directory(root_fd)); + return true; + } + ET_CHECK_OR_RETURN_ERROR( + errno == EEXIST, AccessFailed, "Cannot publish Core AI stored bundle"); + return false; +} + +Error recover_lost_publication_race( + int root_fd, + NSString* key, + NSDictionary* files) { + auto winner = inspect_source_bundle(root_fd, key, files); + if (!winner.ok()) { + return winner.error(); + } + ET_CHECK_OR_RETURN_ERROR( + winner.get() == EntryState::Valid, + InvalidExternalData, + "Concurrently published Core AI source is incomplete"); + return fully_sync_directory(root_fd); +} + +} // namespace + +Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* named_data, + const CacheEntryLock& cache_entry, + CorruptBundlePolicy corrupt_policy) { + NSString* key = cache_entry.key(); + auto prepared = cache_entry.prepare_staging_root(); + if (!prepared.ok()) { + return prepared.error(); + } + NSURL* root = [NSURL fileURLWithPath:prepared.get() isDirectory:YES]; + NSURL* model_url = [[root URLByAppendingPathComponent:key isDirectory:YES] + URLByAppendingPathComponent:bundle.bundle_name + isDirectory:YES]; + FileDescriptor root_fd(retry_eintr([&] { + return open( + root.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + root_fd.get() >= 0, AccessFailed, "Cannot open Core AI source root"); + + auto state = inspect_source_bundle(root_fd.get(), key, bundle.files); + if (!state.ok()) { + return state.error(); + } + if (state.get() == EntryState::Valid) { + ET_CHECK_OK_OR_RETURN_ERROR(fully_sync_directory(root_fd.get())); + return model_url; + } + if (state.get() == EntryState::Corrupt) { + ET_CHECK_OR_RETURN_ERROR( + corrupt_policy == CorruptBundlePolicy::Replace && named_data != nullptr, + InvalidExternalData, + "Cannot replace a damaged Core AI stored bundle"); + ET_CHECK_OK_OR_RETURN_ERROR(remove_storage_directory( + root_fd.get(), key, StorageRemovalPolicy::Maintenance)); + } + ET_CHECK_OR_RETURN_ERROR( + named_data != nullptr, + InvalidProgram, + "Missing NamedDataMap for Core AI source recovery"); + + auto staging_result = create_staging_directory(root_fd.get()); + if (!staging_result.ok()) { + return staging_result.error(); + } + StagingDirectory staging = std::move(staging_result.get()); + auto staging_fd_result = open_staging_directory(root_fd.get(), staging.name); + if (!staging_fd_result.ok()) { + return staging_fd_result.error(); + } + FileDescriptor staging_fd(staging_fd_result.get()); + ET_CHECK_OK_OR_RETURN_ERROR( + write_staging_bundle(staging_fd.get(), bundle, *named_data)); + + auto published = publish_staging_directory(root_fd.get(), staging, key); + if (!published.ok()) { + return published.error(); + } + if (!published.get()) { + ET_CHECK_OK_OR_RETURN_ERROR( + recover_lost_publication_race(root_fd.get(), key, bundle.files)); + } + return model_url; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index 5bcfecfa4f5..33bd34d3272 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -47,4 +47,18 @@ runtime::Result*> storage_children( int fd, StorageNamePolicy policy = StorageNamePolicy::Strict); +enum class StorageRemovalPolicy { + Strict, + Maintenance, +}; + +// Removes one child directory without following links. The complete tree is +// checked before removal and every entry is checked again during removal. +// Missing entries are already removed. Maintenance also removes malformed +// backend-owned non-directory entries. +runtime::Error remove_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy = StorageRemovalPolicy::Strict); + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 066c60e17ba..f30c992081d 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -119,6 +119,62 @@ FileDescriptor file(retry_eintr([&] { } return result; } + +enum class StorageTreePass { + Preflight, + Remove, +}; + +Error walk_storage_tree( + int fd, + StorageTreePass pass, + StorageRemovalPolicy policy) { + auto children = storage_children(fd); + if (!children.ok()) { + return children.error(); + } + for (NSString* name in children.get()) { + struct stat info; + const int inspected = lstat_at(fd, name.fileSystemRepresentation, &info); + if (inspected < 0 && errno == ENOENT) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI storage entry"); + const bool directory = S_ISDIR(info.st_mode); + if (policy == StorageRemovalPolicy::Strict) { + ET_CHECK_OR_RETURN_ERROR( + directory || S_ISREG(info.st_mode), + InvalidExternalData, + "Unexpected Core AI storage entry type"); + } + if (directory) { + FileDescriptor child( + open_directory_at(fd, name.fileSystemRepresentation)); + if (child.get() < 0 && errno == ENOENT) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + child.get() >= 0, + AccessFailed, + "Cannot open Core AI storage directory"); + ET_CHECK_OK_OR_RETURN_ERROR(walk_storage_tree(child.get(), pass, policy)); + } + if (pass == StorageTreePass::Remove) { + const int removed = storage_call(FaultPoint::Remove, [&] { + return unlinkat( + fd, name.fileSystemRepresentation, directory ? AT_REMOVEDIR : 0); + }); + if (removed < 0 && errno == ENOENT) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + removed == 0, AccessFailed, "Cannot remove Core AI storage entry"); + } + } + return pass == StorageTreePass::Remove ? sync_storage_directory(fd) + : Error::Ok; +} } // namespace Result open_storage_directory( @@ -203,6 +259,55 @@ FileDescriptor copy(retry_eintr( return [names sortedArrayUsingSelector:@selector(compare:)]; } +Error remove_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy) { + ET_CHECK_OR_RETURN_ERROR( + storage_component(name), InvalidArgument, "Invalid Core AI entry name"); + struct stat info; + const int inspected = + lstat_at(parent_fd, name.fileSystemRepresentation, &info); + if (inspected < 0 && errno == ENOENT) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI storage entry"); + if (!S_ISDIR(info.st_mode)) { + ET_CHECK_OR_RETURN_ERROR( + policy == StorageRemovalPolicy::Maintenance, + InvalidExternalData, + "Core AI storage entry is not a directory"); + const int removed = storage_call(FaultPoint::Remove, [&] { + return unlinkat(parent_fd, name.fileSystemRepresentation, 0); + }); + ET_CHECK_OR_RETURN_ERROR( + removed == 0 || (removed < 0 && errno == ENOENT), + AccessFailed, + "Cannot remove Core AI storage entry"); + return sync_storage_directory(parent_fd); + } + FileDescriptor entry( + open_directory_at(parent_fd, name.fileSystemRepresentation)); + if (entry.get() < 0 && errno == ENOENT) { + return sync_storage_directory(parent_fd); + } + ET_CHECK_OR_RETURN_ERROR( + entry.get() >= 0, AccessFailed, "Cannot open Core AI storage directory"); + ET_CHECK_OK_OR_RETURN_ERROR( + walk_storage_tree(entry.get(), StorageTreePass::Preflight, policy)); + ET_CHECK_OK_OR_RETURN_ERROR( + walk_storage_tree(entry.get(), StorageTreePass::Remove, policy)); + const int removed = storage_call(FaultPoint::Remove, [&] { + return unlinkat(parent_fd, name.fileSystemRepresentation, AT_REMOVEDIR); + }); + ET_CHECK_OR_RETURN_ERROR( + removed == 0 || (removed < 0 && errno == ENOENT), + AccessFailed, + "Cannot remove Core AI storage directory"); + return sync_storage_directory(parent_fd); +} + int open_storage_shared_file(int parent_fd, const char* name) { // APFS can return ENOENT to a concurrent O_CREAT loser without O_EXCL. int fd = retry_eintr([&] { diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm index 949c5b4ade8..d3a4d46924b 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm @@ -198,6 +198,7 @@ FileDescriptor file(open( } TEST_F(CoreAIBookmarkIOTest, PreparesSharedStagingParentOnDemand) { + EXPECT_TRUE([locked->key() isEqual:key]); NSString* expected = [root.url URLByAppendingPathComponent:@"staging"].path; EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:expected]); auto staging = locked->prepare_staging_root(); diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h index 625f1f48a8a..d15d865020c 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -23,5 +23,12 @@ struct TestDirectory { ::testing::AssertionResult backup_excluded(NSURL* url); ::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded); +::testing::AssertionResult asset_tree_snapshot( + NSURL* root, + NSDictionary* __strong& result); +::testing::AssertionResult snapshot_matches( + NSURL* root, + NSDictionary* expected); +NSArray* staging_directories(NSURL* root); } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm index 6c60dddd6c5..827a228eb8e 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -7,6 +7,7 @@ */ #include "coreai_filesystem_fixture.h" +#include #include namespace executorch::backends::coreai::testing { @@ -68,4 +69,101 @@ << "Backup flag did not settle: " << url.path.UTF8String; } +::testing::AssertionResult asset_tree_snapshot( + NSURL* root, + NSDictionary* __strong& result) { + result = nil; + if (root == nil) + return ::testing::AssertionFailure() << "Missing snapshot root"; + NSFileManager* manager = NSFileManager.defaultManager; + NSError* error = nil; + NSArray* children = [manager subpathsOfDirectoryAtPath:root.path + error:&error]; + if (children == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot list " << root.path.UTF8String; + } + NSMutableDictionary* snapshot = [NSMutableDictionary dictionary]; + for (NSString* path in [@[ @"" ] arrayByAddingObjectsFromArray:children]) { + NSURL* url = + path.length == 0 ? root : [root URLByAppendingPathComponent:path]; + struct stat info; + if (lstat(url.fileSystemRepresentation, &info) != 0) { + return ::testing::AssertionFailure() + << "Cannot stat " << url.path.UTF8String; + } + NSMutableDictionary* item = [@{ + @"inode" : @(info.st_ino), + @"device" : @(info.st_dev), + @"mode" : @(info.st_mode), + @"size" : @(info.st_size), + @"mtime_seconds" : @(info.st_mtimespec.tv_sec), + @"mtime_nanos" : @(info.st_mtimespec.tv_nsec) + } mutableCopy]; + error = nil; + if (S_ISREG(info.st_mode)) { + NSData* bytes = [NSData dataWithContentsOfURL:url options:0 error:&error]; + if (bytes == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot read " << url.path.UTF8String; + } + item[@"bytes"] = bytes; + } else if (S_ISLNK(info.st_mode)) { + NSString* destination = [manager destinationOfSymbolicLinkAtPath:url.path + error:&error]; + if (destination == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot read link " << url.path.UTF8String; + } + item[@"link"] = destination; + } + snapshot[path] = item; + } + result = snapshot; + return ::testing::AssertionSuccess(); +} + +::testing::AssertionResult snapshot_matches( + NSURL* root, + NSDictionary* expected) { + if (expected == nil) + return ::testing::AssertionFailure() << "Missing expected snapshot"; + NSDictionary* actual = nil; + auto read = asset_tree_snapshot(root, actual); + if (!read) + return read; + return [actual isEqual:expected] ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() + << "Filesystem snapshot changed: " << root.path.UTF8String; +} + +NSArray* staging_directories(NSURL* root) { + if (root == nil) { + ADD_FAILURE() << "Missing staging root"; + return nil; + } + NSURL* cache_staging = [root URLByAppendingPathComponent:@"staging" + isDirectory:YES]; + if ([NSFileManager.defaultManager fileExistsAtPath:cache_staging.path]) { + root = cache_staging; + } + NSError* error = nil; + NSArray* children = + [NSFileManager.defaultManager contentsOfDirectoryAtURL:root + includingPropertiesForKeys:nil + options:0 + error:&error]; + if (children == nil || error != nil) { + ADD_FAILURE() << "Cannot list staging directories: " + << root.path.UTF8String; + return nil; + } + NSMutableArray* staging = [NSMutableArray array]; + for (NSURL* child in children) { + if ([child.lastPathComponent hasPrefix:@".staging-"]) + [staging addObject:child]; + } + return staging; +} + } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.h b/backends/apple/coreai/runtime/test/coreai_source_fixture.h new file mode 100644 index 00000000000..561466eb208 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include "coreai_filesystem_fixture.h" +#include "coreai_manifest_fixture.h" +#include "coreai_source_bundle.h" + +namespace executorch::backends::coreai::testing { + +class TestData final : public runtime::NamedDataMap { + public: + mutable std::atomic releases{0}; + mutable std::atomic requests{0}; + mutable std::atomic attempts{0}; + mutable std::atomic metadata_requests{0}; + int fail_at = 0; + runtime::Error failure_error = runtime::Error::AccessFailed; + std::string required_prefix; + std::map files{ + {"coreai/ab/model.aimodel/graph.bin", std::string("model data", 11)}}; + runtime::Result get_tensor_layout( + std::string_view) const override; + runtime::Result get_data( + std::string_view key) const override; + runtime::Error load_data_into(std::string_view, void*, size_t) const override; + runtime::Result get_num_keys() const override; + runtime::Result get_key(uint32_t i) const override; +}; + +void aot_data(TestData& data, NSString* arch); + +class CoreAISourceTest : public ::testing::Test { + protected: + void check_concurrent_source(); + runtime::Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_root, + CorruptBundlePolicy corrupt_policy = CorruptBundlePolicy::Reject); +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.mm b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm new file mode 100644 index 00000000000..fbed96d3fa4 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm @@ -0,0 +1,96 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "coreai_source_fixture.h" + +#include +#include "coreai_bookmarks.h" +#include "coreai_cache_key.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +Result TestData::get_tensor_layout( + std::string_view) const { + return Error::NotSupported; +} + +Result TestData::get_data(std::string_view key) const { + if (!required_prefix.empty() && + key.substr(0, required_prefix.size()) != required_prefix) { + ADD_FAILURE() << "NamedDataMap key " << key + << " is outside selected prefix " << required_prefix; + return Error::InvalidArgument; + } + if (++attempts == fail_at) { + return failure_error; + } + auto found = files.find(std::string(key)); + if (found == files.end()) { + return Error::NotFound; + } + ++requests; + return runtime::FreeableBuffer( + found->second.data(), + found->second.size(), + [](void* context, void*, size_t) { + ++static_cast(context)->releases; + }, + const_cast(this)); +} + +Error TestData::load_data_into(std::string_view, void*, size_t) const { + return Error::NotSupported; +} + +Result TestData::get_num_keys() const { + ++metadata_requests; + return static_cast(files.size()); +} + +Result TestData::get_key(uint32_t i) const { + ++metadata_requests; + if (i >= files.size()) { + return Error::NotFound; + } + auto found = files.begin(); + std::advance(found, i); + return found->first.c_str(); +} + +void aot_data(TestData& data, NSString* arch) { + data.required_prefix = + [NSString stringWithFormat:@"coreai/ab/model.%@.aimodelc/", arch] + .UTF8String; + data.files.clear(); + data.files[data.required_prefix + "graph.bin"] = + std::string("compiled\0graph", 14); + data.files[data.required_prefix + "nested/weights.bin"] = "weights"; +} + +Result CoreAISourceTest::prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_root, + CorruptBundlePolicy corrupt_policy) { + auto key = cache_variant_key(bundle, @"macOS", @"arch_b", @"default"); + if (!key.ok()) { + ADD_FAILURE() << "Cannot compute source fixture key"; + return key.error(); + } + auto lock = lock_cache_entry(cache_root, key.get()); + if (!lock.ok()) { + ADD_FAILURE() << "Cannot acquire source fixture lock"; + return lock.error(); + } + return executorch::backends::coreai::prepare_source_bundle( + bundle, data, *lock.get(), corrupt_policy); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_test.mm b/backends/apple/coreai/runtime/test/coreai_source_test.mm new file mode 100644 index 00000000000..301dafad56e --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_test.mm @@ -0,0 +1,979 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include "coreai_bookmarks.h" +#include "coreai_cache_key.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +namespace { +Result selected_bundle( + NSDictionary* dictionary, + NSString* architecture = @"arch_b") { + auto parsed = parse_manifest(encode(dictionary)); + if (!parsed.ok()) { + return parsed.error(); + } + return select_assets(parsed.get(), architecture, @"macOS"); +} + +Result source_manifest(bool aot = false) { + return selected_bundle(aot ? aot_manifest_dict() : manifest_dict()); +} + +NSString* source_key(const SelectedBundle& bundle) { + auto key = cache_variant_key(bundle, @"macOS", @"arch_b", @"default"); + EXPECT_TRUE(key.ok()); + return key.ok() ? key.get() : nil; +} + +NSURL* staging_root(NSURL* cache_root) { + return [cache_root URLByAppendingPathComponent:@"staging" isDirectory:YES]; +} + +NSURL* source_entry(NSURL* cache_root, const SelectedBundle& bundle) { + return + [staging_root(cache_root) URLByAppendingPathComponent:source_key(bundle) + isDirectory:YES]; +} +} // namespace + +TEST_F(CoreAISourceTest, UnicodePathsMaterializeAndReuse) { + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + for (NSString* name in + @[ composed, composed.decomposedStringWithCanonicalMapping ]) { + SCOPED_TRACE(name.UTF8String); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files.clear(); + NSString* relative = [NSString stringWithFormat:@"%@/%@.bin", name, name]; + NSString* file = [@"model.aimodel" stringByAppendingPathComponent:relative]; + auto dict = manifest_dict(); + dict[@"files"] = @{file : @7}; + data.files[[NSString stringWithFormat:@"coreai/ab/%@", file].UTF8String] = + "unicode"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto source = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSData* bytes = [NSData + dataWithContentsOfURL:[source.get() + URLByAppendingPathComponent:relative]]; + EXPECT_TRUE([bytes isEqual:[NSData dataWithBytes:"unicode" length:7]]); + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:source.get()]); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + } +} + +TEST_F(CoreAISourceTest, RejectsAliasedPathsBeforeReadingAssets) { + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + NSString* decomposed = composed.decomposedStringWithCanonicalMapping; + for (NSArray* pair in @[ + @[ @"A.bin", @"a.bin" ], + @[ @"Dir/x", @"dir/y" ], + @[ @"Dir", @"dir/x" ], + @[ composed, decomposed ], + @[ + [composed stringByAppendingPathComponent:@"x"], + [decomposed stringByAppendingPathComponent:@"y"] + ] + ]) { + SCOPED_TRACE(pair.description.UTF8String); + TestData data; + auto dictionary = manifest_dict(); + dictionary[@"files"] = @{ + [@"model.aimodel" stringByAppendingPathComponent:pair[0]] : @1, + [@"model.aimodel" stringByAppendingPathComponent:pair[1]] : @1 + }; + EXPECT_EQ( + parse_manifest(encode(dictionary)).error(), Error::InvalidProgram); + EXPECT_EQ(data.attempts, 0); + } +} + +TEST_F(CoreAISourceTest, AotMaterializesOnlySelectedArchitecture) { + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { + SCOPED_TRACE(arch.UTF8String); + TestData data; + aot_data(data, arch); + auto selected = select_assets(parsed.get(), arch, @"macOS"); + ASSERT_TRUE(selected.ok()); + auto first = prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.requests, 2); + NSURL* url = first.get(); + EXPECT_TRUE([url.lastPathComponent + isEqualToString:selected->path.lastPathComponent]); + NSData* contents = [NSData + dataWithContentsOfURL:[url URLByAppendingPathComponent:@"graph.bin"]]; + ASSERT_TRUE([contents isEqual:[NSData dataWithBytes:"compiled\0graph" + length:14]]); + + NSString* other_arch = + [arch isEqualToString:@"arch_a"] ? @"arch_b" : @"arch_a"; + NSString* other_key = [NSString + stringWithFormat:@"coreai/ab/model.%@.aimodelc/graph.bin", other_arch]; + data.files[other_key.UTF8String] = "must not affect identity"; + auto reused = + prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + + auto unselected_changed = aot_manifest_dict(); + NSMutableDictionary* digests = + [unselected_changed[@"bundle_digests"] mutableCopy]; + digests[[NSString stringWithFormat:@"model.%@.aimodelc", other_arch]] = + fixture_identity('f'); + unselected_changed[@"bundle_digests"] = digests; + NSMutableDictionary* files = [unselected_changed[@"files"] mutableCopy]; + files[[NSString stringWithFormat:@"model.%@.aimodelc/extra", other_arch]] = + @0; + unselected_changed[@"files"] = files; + auto changed_parsed = parse_manifest(encode(unselected_changed)); + ASSERT_TRUE(changed_parsed.ok()); + auto unchanged = select_assets(changed_parsed.get(), arch, @"macOS"); + ASSERT_TRUE(unchanged.ok()); + auto unchanged_source = + prepare_source_bundle(unchanged.get(), &data, storage.url.path); + ASSERT_TRUE(unchanged_source.ok()); + EXPECT_TRUE([unchanged_source.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.releases, data.requests); + } +} + +void CoreAISourceTest::check_concurrent_source() { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto selected = source_manifest(); + ASSERT_TRUE(selected.ok()); + for (int pass = 0; pass < 2; ++pass) { + SCOPED_TRACE(pass == 0 ? "cold" : "warm"); + using FixturePointer = decltype(this); + struct Worker { + FixturePointer fixture = nullptr; + const SelectedBundle* manifest = nullptr; + TestData* data = nullptr; + NSString* root = nil; + dispatch_semaphore_t start = nullptr; + bool started = false; + Error error = Error::Internal; + std::string path; + } workers[4]; + pthread_t threads[4]; + dispatch_semaphore_t start = dispatch_semaphore_create(0); + int launch_error = 0; + size_t launched = 0; + for (; launched < 4; ++launched) { + auto& worker = workers[launched]; + worker.fixture = this; + worker.manifest = &selected.get(); + worker.data = &data; + worker.root = storage.url.path; + worker.start = start; + launch_error = pthread_create( + &threads[launched], + nullptr, + [](void* context) -> void* { + auto& worker = *static_cast(context); + @autoreleasepool { + worker.started = + dispatch_semaphore_wait( + worker.start, + dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)) == 0; + if (!worker.started) + return nullptr; + auto assets = worker.fixture->prepare_source_bundle( + *worker.manifest, worker.data, worker.root); + worker.error = assets.error(); + if (assets.ok()) + worker.path = assets.get().path.UTF8String; + } + return nullptr; + }, + &worker); + if (launch_error != 0) + break; + } + for (size_t i = 0; i < launched; ++i) + dispatch_semaphore_signal(start); + int join_errors[4] = {}; + for (size_t i = 0; i < launched; ++i) + join_errors[i] = pthread_join(threads[i], nullptr); + ASSERT_EQ(launch_error, 0); + for (size_t i = 0; i < 4; ++i) { + SCOPED_TRACE(i); + ASSERT_EQ(join_errors[i], 0); + ASSERT_TRUE(workers[i].started); + ASSERT_EQ(workers[i].error, Error::Ok); + EXPECT_EQ(workers[i].path, workers[0].path); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:staging_root(storage.url).path + error:nil] + .count, + 1u); + } +} + +TEST_F(CoreAISourceTest, PerKeyLockSerializesColdAndWarmPreparation) { + check_concurrent_source(); +} + +TEST_F(CoreAISourceTest, MissingSourceRequiresNamedData) { + auto bundle = source_manifest(); + ASSERT_TRUE(bundle.ok()); + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + EXPECT_EQ( + prepare_source_bundle(bundle.get(), nullptr, storage.url.path).error(), + Error::InvalidProgram); + auto recovered = prepare_source_bundle(bundle.get(), &data, storage.url.path); + ASSERT_TRUE(recovered.ok()); + EXPECT_EQ( + [NSData + dataWithContentsOfURL:[recovered.get() + URLByAppendingPathComponent:@"graph.bin"]] + .length, + 11u); + EXPECT_EQ(data.requests, data.releases); +} + +TEST_F(CoreAISourceTest, PublishedSourceOutlivesResult) { + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSURL* first_url = nil; + { + auto assets = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(assets.ok()); + EXPECT_EQ(data.releases, data.requests); + first_url = assets.get(); + EXPECT_EQ( + [NSData + dataWithContentsOfURL:[first_url + URLByAppendingPathComponent:@"graph.bin"]] + .length, + 11u); + } + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:first_url.path]); + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first_url]); +} + +TEST_F(CoreAISourceTest, PersistentIdentityTracksDigestNotFileOrder) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/weights.bin" : @7 + }; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + NSURL* first_url = first.get(); + dict[@"files"] = @{ + @"model.aimodel/nested/weights.bin" : @7, + @"model.aimodel/graph.bin" : @11 + }; + auto reordered_manifest = selected_bundle(dict); + ASSERT_TRUE(reordered_manifest.ok()); + auto reordered = + prepare_source_bundle(reordered_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(reordered.ok()); + EXPECT_TRUE([reordered.get() isEqual:first_url]); + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "changed"; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('3')}; + auto changed_manifest = selected_bundle(dict); + ASSERT_TRUE(changed_manifest.ok()); + auto changed = + prepare_source_bundle(changed_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(changed.ok()); + EXPECT_FALSE([changed.get() isEqual:first_url]); + EXPECT_TRUE( + [[NSData dataWithContentsOfURL:[first_url URLByAppendingPathComponent: + @"nested/weights.bin"]] + isEqual:[@"weights" dataUsingEncoding:NSUTF8StringEncoding]]); + EXPECT_EQ(data.releases, data.requests); +} + +TEST_F(CoreAISourceTest, PartialMaterializationCleansStagingAndCanRetry) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/weights.bin" : @7 + }; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + data.fail_at = 2; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::AccessFailed); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:staging_root(storage.url).path + error:nil] + .count, + 0u); + EXPECT_EQ(data.releases, data.requests); + data.fail_at = 0; + EXPECT_TRUE( + prepare_source_bundle(manifest.get(), &data, storage.url.path).ok()); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:staging_root(storage.url).path + error:nil] + .count, + 1u); +} + +TEST_F( + CoreAISourceTest, + NestedEmptyAndLargeFilesMaterializeAndReuseWithoutNamedData) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + constexpr size_t large_size = 2 * 1024 * 1024 + 17; + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/empty" : @0, + @"model.aimodel/nested/deeper/large.bin" : @(large_size) + }; + data.files["coreai/ab/model.aimodel/nested/empty"] = ""; + data.files["coreai/ab/model.aimodel/nested/deeper/large.bin"] = + std::string(large_size, 'x'); + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 3); + EXPECT_EQ(data.requests, data.releases); + NSData* empty = [NSData + dataWithContentsOfURL:[first.get() + URLByAppendingPathComponent:@"nested/empty"]]; + ASSERT_NE(empty, nil); + EXPECT_EQ(empty.length, 0u); + NSData* large = [NSData + dataWithContentsOfURL: + [first.get() URLByAppendingPathComponent:@"nested/deeper/large.bin"]]; + ASSERT_NE(large, nil); + EXPECT_EQ(large.length, large_size); + + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + const int attempts = data.attempts; + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); +} + +TEST_F( + CoreAISourceTest, + ExplicitStagingKeyReusesSourceWithoutWritesOrRepublication) { + TestDirectory directory; + ASSERT_NE(directory.url, nil); + TestData data; + auto bundle = source_manifest(); + ASSERT_TRUE(bundle.ok()); + auto key = cache_variant_key( + bundle.get(), @"macOS", @"arch_b", @"explicit-test-options"); + ASSERT_TRUE(key.ok()); + auto lock = lock_cache_entry(directory.url.path, key.get()); + ASSERT_TRUE(lock.ok()); + auto first = executorch::backends::coreai::prepare_source_bundle( + bundle.get(), &data, *lock.get()); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 1); + NSString* expected_path = [directory.url.path + stringByAppendingPathComponent: + [NSString stringWithFormat:@"staging/%@/model.aimodel", key.get()]]; + EXPECT_TRUE([first.get().path isEqualToString:expected_path]); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(directory.url, snapshot)); + const int attempts = data.attempts; + const int requests = data.requests; + data.files.clear(); + auto reused = executorch::backends::coreai::prepare_source_bundle( + bundle.get(), &data, *lock.get()); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, requests); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(directory.url, snapshot)); +} + +TEST_F(CoreAISourceTest, RejectsMalformedStoredBundlesWithoutReplacement) { + for (bool aot : {false, true}) { + SCOPED_TRACE(aot ? "aot" : "source"); + for (NSString* damage in @[ + @"missing-file", + @"extra-file", + @"wrong-size", + @"symlink-file", + @"file-at-key", + @"extra-directory" + ]) { + SCOPED_TRACE(damage.UTF8String); + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + TestData data; + if (aot) + aot_data(data, @"arch_b"); + auto manifest = source_manifest(aot); + ASSERT_TRUE(manifest.ok()); + auto source = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSURL* entry = source.get().URLByDeletingLastPathComponent; + NSURL* graph = [source.get() URLByAppendingPathComponent:@"graph.bin"]; + NSData* bytes = [NSData dataWithContentsOfURL:graph]; + ASSERT_NE(bytes, nil); + NSFileManager* manager = NSFileManager.defaultManager; + if ([damage isEqual:@"missing-file"]) { + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + } else if ([damage isEqual:@"extra-file"]) { + ASSERT_TRUE([[NSData data] + writeToURL:[source.get() URLByAppendingPathComponent:@"extra"] + atomically:NO]); + } else if ([damage isEqual:@"wrong-size"]) { + ASSERT_TRUE([[NSData data] writeToURL:graph atomically:NO]); + } else if ([damage isEqual:@"symlink-file"]) { + NSURL* destination = + [target.url URLByAppendingPathComponent:@"graph.bin"]; + ASSERT_TRUE([bytes writeToURL:destination atomically:NO]); + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + ASSERT_TRUE([manager createSymbolicLinkAtPath:graph.path + withDestinationPath:destination.path + error:nil]); + } else if ([damage isEqual:@"file-at-key"]) { + ASSERT_TRUE([manager removeItemAtURL:entry error:nil]); + ASSERT_TRUE([bytes writeToURL:entry atomically:NO]); + } else { + ASSERT_TRUE([manager + createDirectoryAtURL: + [source.get() URLByAppendingPathComponent:@"extra"] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + } + NSDictionary* snapshot = nil; + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + const int attempts = data.attempts; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path) + .error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, attempts); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + } + } +} + +TEST_F(CoreAISourceTest, ExplicitReplacementRepairsMalformedStoredBundles) { + for (NSString* damage in @[ + @"missing-file", + @"extra-file", + @"wrong-size", + @"symlink-file", + @"fifo-file", + @"file-at-key", + @"symlink-at-key", + @"extra-directory" + ]) { + SCOPED_TRACE(damage.UTF8String); + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + TestData data; + auto bundle = source_manifest(); + ASSERT_TRUE(bundle.ok()); + auto source = prepare_source_bundle(bundle.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSURL* entry = source.get().URLByDeletingLastPathComponent; + NSURL* graph = [source.get() URLByAppendingPathComponent:@"graph.bin"]; + NSFileManager* manager = NSFileManager.defaultManager; + if ([damage isEqual:@"file-at-key"] || [damage isEqual:@"symlink-at-key"]) { + ASSERT_TRUE([manager removeItemAtURL:entry error:nil]); + if ([damage isEqual:@"file-at-key"]) { + ASSERT_TRUE([[NSData data] writeToURL:entry atomically:NO]); + } else { + ASSERT_EQ( + symlink( + target.url.fileSystemRepresentation, + entry.fileSystemRepresentation), + 0); + } + } else if ([damage isEqual:@"extra-file"]) { + ASSERT_TRUE([[NSData data] + writeToURL:[source.get() URLByAppendingPathComponent:@"extra"] + atomically:NO]); + } else if ([damage isEqual:@"extra-directory"]) { + ASSERT_TRUE([manager + createDirectoryAtURL:[source.get() + URLByAppendingPathComponent:@"extra"] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + } else { + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + if ([damage isEqual:@"wrong-size"]) { + ASSERT_TRUE([[NSData data] writeToURL:graph atomically:NO]); + } else if ([damage isEqual:@"symlink-file"]) { + NSURL* sentinel = [target.url URLByAppendingPathComponent:@"sentinel"]; + ASSERT_TRUE([@"external" writeToURL:sentinel + atomically:NO + encoding:NSUTF8StringEncoding + error:nil]); + ASSERT_EQ( + symlink( + sentinel.fileSystemRepresentation, + graph.fileSystemRepresentation), + 0); + } else if ([damage isEqual:@"fifo-file"]) { + ASSERT_EQ(mkfifo(graph.fileSystemRepresentation, 0600), 0); + } + } + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + auto repaired = prepare_source_bundle( + bundle.get(), &data, storage.url.path, CorruptBundlePolicy::Replace); + ASSERT_TRUE(repaired.ok()); + EXPECT_TRUE([repaired.get() isEqual:source.get()]); + EXPECT_TRUE(([[NSData + dataWithContentsOfURL:[repaired.get() + URLByAppendingPathComponent:@"graph.bin"]] + isEqual:[NSData dataWithBytes:"model data" length:11]])); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + } +} + +TEST_F(CoreAISourceTest, PublishedDirectoriesUseDefaultPermissions) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* reference = [storage.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + for (NSURL* directory in + @[ source.get().URLByDeletingLastPathComponent, source.get() ]) { + SCOPED_TRACE(directory.path.UTF8String); + struct stat actual; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST_F(CoreAISourceTest, PermissionErrorsRemainDistinctFromCorruptEntries) { + if (geteuid() == 0) { + GTEST_SKIP() << "Root can bypass directory read permissions"; + } + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + for (NSURL* directory in + @[ source.get().URLByDeletingLastPathComponent, source.get() ]) { + SCOPED_TRACE(directory.path.UTF8String); + struct stat original; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &original), 0); + ASSERT_EQ(chmod(directory.fileSystemRepresentation, 0000), 0); + auto denied = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + EXPECT_EQ( + chmod(directory.fileSystemRepresentation, original.st_mode & 07777), 0); + EXPECT_EQ(denied.error(), Error::AccessFailed); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_EQ(data.attempts, 1); + EXPECT_TRUE( + prepare_source_bundle(manifest.get(), nullptr, storage.url.path).ok()); + } +} + +TEST_F(CoreAISourceTest, PreservesNamedDataErrorsWithoutRetrying) { + for (Error error : {Error::AccessFailed, Error::MemoryAllocationFailed}) { + SCOPED_TRACE(static_cast(error)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.fail_at = 1; + data.failure_error = error; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + error); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 0); + EXPECT_EQ(data.releases, 0); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:source_entry(storage.url, manifest.get()).path]); + } +} + +TEST_F(CoreAISourceTest, MissingNamedDataIsInvalidExternalData) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files.clear(); + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 0); + EXPECT_EQ(staging_directories(storage.url).count, 0u); +} + +TEST_F(CoreAISourceTest, RejectsMismatchedBufferSizesAndReleasesPayload) { + for (size_t size : {0u, 10u, 12u}) { + SCOPED_TRACE(size); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files["coreai/ab/model.aimodel/graph.bin"] = std::string(size, 'x'); + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 1); + EXPECT_EQ(data.releases, 1); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:source_entry(storage.url, manifest.get()).path]); + } +} + +TEST_F(CoreAISourceTest, FullyFlushesSourceBeforeAndAfterPublication) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + aot_data(data, @"arch_b"); + auto manifest = source_manifest(true); + ASSERT_TRUE(manifest.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + NSURL* root = staging_root(storage.url); + const NSUInteger file_count = manifest->files.count; + __block int synced_files = 0; + __block int full_syncs = 0; + __block int publications = 0; + __block bool ordered = true; + StorageFaultScope fault(FaultPoint::FullSync); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::SyncFile) + ++synced_files; + if (operation == FaultPoint::FullSync) { + ++full_syncs; + ordered &= synced_files == file_count; + const bool published = + [NSFileManager.defaultManager fileExistsAtPath:final.path]; + ordered &= full_syncs == 1 ? !published : published; + ordered &= staging_directories(root).count == (full_syncs == 1 ? 1u : 0u); + } + if (operation == FaultPoint::Rename) { + ++publications; + ordered &= full_syncs == 1; + } + }; + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + EXPECT_TRUE(ordered); + EXPECT_EQ(full_syncs, 2); + EXPECT_EQ(publications, 1); + EXPECT_EQ(data.attempts, file_count); + EXPECT_EQ(data.requests, data.releases); +} + +TEST_F(CoreAISourceTest, PublicationFailuresCleanStagingWithoutAutomaticRetry) { + for (auto operation : + {FaultPoint::CreateStagingDirectory, + FaultPoint::OpenStagingDirectory, + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::Rename}) { + SCOPED_TRACE(static_cast(operation)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSString* key = source_key(manifest.get()); + auto lock = lock_cache_entry(storage.url.path, key); + ASSERT_TRUE(lock.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + { + StorageFaultScope fault(operation); + EXPECT_EQ( + executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()) + .error(), + Error::AccessFailed); + EXPECT_EQ(fault.matches, 1); + EXPECT_EQ(fault.hits, 1); + } + const int attempts = operation == FaultPoint::CreateStagingDirectory || + operation == FaultPoint::OpenStagingDirectory || + operation == FaultPoint::SyncDirectory + ? 0 + : 1; + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, data.releases); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:final.path]); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + auto retried = executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()); + ASSERT_TRUE(retried.ok()); + EXPECT_EQ(data.attempts, attempts + 1); + } +} + +TEST_F(CoreAISourceTest, RetriesInterruptedStagingAndPublicationCalls) { + for (auto operation : + {FaultPoint::CreateStagingDirectory, + FaultPoint::OpenStagingDirectory, + FaultPoint::SyncDirectory, + FaultPoint::Rename}) { + SCOPED_TRACE(static_cast(operation)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSString* key = source_key(manifest.get()); + auto lock = lock_cache_entry(storage.url.path, key); + ASSERT_TRUE(lock.ok()); + StorageFaultScope fault(operation); + fault.error = EINTR; + auto source = executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()); + ASSERT_TRUE(source.ok()); + EXPECT_EQ(fault.hits, 1); + EXPECT_GE(fault.matches, 2); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + } +} + +TEST_F(CoreAISourceTest, PostRenameFlushFailureLeavesCompleteSourceVisible) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + { + StorageFaultScope fault(FaultPoint::FullSync); + fault.match_index = 2; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::AccessFailed); + EXPECT_EQ(fault.matches, 2); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + NSURL* graph = [final URLByAppendingPathComponent:@"model.aimodel/graph.bin"]; + EXPECT_TRUE(([[NSData dataWithContentsOfURL:graph] + isEqual:[NSData dataWithBytes:"model data" length:11]])); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + { + StorageFaultScope fault(FaultPoint::FullSync); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), nullptr, storage.url.path) + .error(), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + } + StorageFaultScope flush(FaultPoint::FullSync); + flush.error = EINTR; + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_EQ(flush.hits, 1); + EXPECT_EQ(flush.matches, 2); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_EQ(data.attempts, 1); +} + +TEST_F(CoreAISourceTest, IncompleteConcurrentWinnerFailsWithoutReplacement) { + TestDirectory storage; + TestDirectory winner_storage; + ASSERT_NE(storage.url, nil); + ASSERT_NE(winner_storage.url, nil); + TestData data; + TestData winner_data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto winner = prepare_source_bundle( + manifest.get(), &winner_data, winner_storage.url.path); + ASSERT_TRUE(winner.ok()); + ASSERT_TRUE([NSFileManager.defaultManager + removeItemAtURL:[winner.get() URLByAppendingPathComponent:@"graph.bin"] + error:nil]); + NSURL* winner_entry = winner.get().URLByDeletingLastPathComponent; + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); + NSURL* root = staging_root(storage.url); + NSURL* final = source_entry(storage.url, manifest.get()); + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(FaultPoint::Rename); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Rename) + return; + ++publications; + callback_ok &= staging_directories(root).count == 1; + callback_ok &= [NSFileManager.defaultManager moveItemAtURL:winner_entry + toURL:final + error:nil]; + }; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_TRUE(callback_ok); + EXPECT_EQ(publications, 1); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(final, snapshot)); +} + +TEST_F( + CoreAISourceTest, + ExclusivePublicationValidatesWinnerAndRemovesLosingStaging) { + TestDirectory storage; + TestDirectory winner_storage; + ASSERT_NE(storage.url, nil); + ASSERT_NE(winner_storage.url, nil); + TestData data; + TestData winner_data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto winner = prepare_source_bundle( + manifest.get(), &winner_data, winner_storage.url.path); + ASSERT_TRUE(winner.ok()); + NSURL* winner_entry = winner.get().URLByDeletingLastPathComponent; + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); + NSURL* root = staging_root(storage.url); + NSURL* final = source_entry(storage.url, manifest.get()); + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(FaultPoint::FullSync); + fault.match_index = 2; + fault.error = EINTR; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Rename) + return; + ++publications; + callback_ok &= staging_directories(root).count == 1; + callback_ok &= [NSFileManager.defaultManager moveItemAtURL:winner_entry + toURL:final + error:nil]; + }; + auto result = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + EXPECT_TRUE(callback_ok); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE([result.get() + isEqual:[final URLByAppendingPathComponent:manifest->bundle_name]]); + EXPECT_EQ(publications, 1); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(fault.matches, 3); + } + EXPECT_EQ(data.attempts, manifest->files.count); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(final, snapshot)); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 266dfbd34bf..c169197fc01 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -398,6 +398,162 @@ FileDescriptor root(open( EXPECT_EQ(lseek(created.get(), 0, SEEK_CUR), 4); } +TEST(CoreAIStorageTest, StrictRemovalPreflightsBeforeMutation) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a/file", "x", 1), Error::Ok); + ASSERT_EQ( + symlinkat(target.url.fileSystemRepresentation, entry.get(), "z-link"), 0); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot( + [storage.url URLByAppendingPathComponent:@"entry"], snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::InvalidExternalData); + EXPECT_TRUE(snapshot_matches( + [storage.url URLByAppendingPathComponent:@"entry"], snapshot)); +} + +TEST(CoreAIStorageTest, StrictRemovalRevalidatesDuringRemoval) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a", "x", 1), Error::Ok); + ASSERT_EQ(write_storage_file(entry.get(), @"z", "x", 1), Error::Ok); + __block int removals = 0; + const int entry_fd = entry.get(); + NSString* target_path = target.url.path; + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Remove || ++removals != 1) { + return; + } + EXPECT_EQ(unlinkat(entry_fd, "z", 0), 0); + EXPECT_EQ( + symlinkat(target_path.fileSystemRepresentation, entry_fd, "z"), 0); + }; + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::InvalidExternalData); + EXPECT_EQ(removals, 1); + struct stat info; + EXPECT_EQ(lstat_at(entry.get(), "a", &info), -1); + ASSERT_EQ(lstat_at(entry.get(), "z", &info), 0); + EXPECT_TRUE(S_ISLNK(info.st_mode)); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:target.url.path]); +} + +TEST(CoreAIStorageTest, MaintenanceRemovalDoesNotFollowMalformedEntries) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"nested/file", "x", 1), Error::Ok); + ASSERT_EQ( + symlinkat(target.url.fileSystemRepresentation, entry.get(), "link"), 0); + ASSERT_EQ(mkfifoat(entry.get(), "fifo", 0600), 0); + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Maintenance), + Error::Ok); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Maintenance), + Error::Ok); +} + +TEST(CoreAIStorageTest, EntriesMissingDuringRemovalAreAlreadyRemoved) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a", "x", 1), Error::Ok); + ASSERT_EQ(write_storage_file(entry.get(), @"z", "x", 1), Error::Ok); + const int entry_fd = entry.get(); + __block int removals = 0; + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::Remove && ++removals == 1) { + EXPECT_EQ(unlinkat(entry_fd, "z", 0), 0); + } + }; + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::Ok); + EXPECT_EQ(removals, 2); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); +} + +TEST(CoreAIStorageTest, RemovalFailuresCanBeRetried) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"nested/file", "x", 1), Error::Ok); + { + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::Ok); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); +} + TEST(CoreAIStorageTest, RetriesEintr) { int calls = 0; const ssize_t count = retry_eintr([&]() -> ssize_t {