diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index c8b05189631..baaae9a50a9 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -37,7 +37,9 @@ if(EXECUTORCH_BUILD_TESTS) runtime/test/coreai_manifest_test.mm runtime/test/coreai_storage_test.mm runtime/test/coreai_filesystem_fixture.mm + runtime/test/coreai_source_test.mm runtime/test/coreai_bookmark_test.mm + runtime/test/coreai_source_fixture.mm runtime/test/coreai_bookmark_fixture.mm ${_coreai_runtime_sources} ) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 09534810ced..a53e6eb1709 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -51,6 +51,29 @@ processes of the same user rebinding paths. Written files and changed directorie are synced with `fsync`; atomic publication also uses `F_FULLFSYNC` before its rename. +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. +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 +loading without replacement, since an SDK model may still be using them. + +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 a control for iCloud Drive synchronization. diff --git a/backends/apple/coreai/runtime/coreai_assets.h b/backends/apple/coreai/runtime/coreai_assets.h index ebf84a3bf01..4a60dc411ee 100644 --- a/backends/apple/coreai/runtime/coreai_assets.h +++ b/backends/apple/coreai/runtime/coreai_assets.h @@ -9,6 +9,7 @@ #pragma once #import +#include #include namespace executorch::backends::coreai { @@ -33,5 +34,12 @@ runtime::Result select_assets( const Manifest& manifest, NSString* device_architecture, NSString* platform); +// Caller holds the key's disk lock and has prepared staging_root. +// Inline bundles are atomically published at staging_root/key/bundle. +runtime::Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* named_data, + NSString* staging_root, + NSString* key); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_assets.mm b/backends/apple/coreai/runtime/coreai_assets.mm index d28cb8e1d55..f2b547b9217 100644 --- a/backends/apple/coreai/runtime/coreai_assets.mm +++ b/backends/apple/coreai/runtime/coreai_assets.mm @@ -7,14 +7,29 @@ */ #import "coreai_assets.h" +#include "coreai_file.h" + +#include +#include +#include +#include +#include #include +#include #include +#import "coreai_storage.h" namespace executorch::backends::coreai { namespace { using runtime::Error; using runtime::Result; +Error path_error(int error) { + return error == ENOENT || error == ENOTDIR || error == ELOOP + ? Error::InvalidExternalData + : Error::AccessFailed; +} + bool nonempty_string(id value) { return [value isKindOfClass:NSString.class] && [value length] > 0 && [value rangeOfString:@"\0"].location == NSNotFound; @@ -121,6 +136,111 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return numbers[0] > 0; } +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]; + struct stat info; + if (retry_eintr([&] { + return fstatat(fd, name.fileSystemRepresentation, &info, + AT_SYMLINK_NOFOLLOW); + }) != 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(retry_eintr([&] { + return openat(fd, name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + 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; +} + +// Checks the exact file set, entry types and sizes. Contents are not hashed. +Error validate_tree(int fd, NSDictionary* 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; +} + +enum class EntryState { Missing, Valid, Corrupt }; + +Result inspect_entry( + int root_fd, + NSString* name, + NSDictionary* files) { + FileDescriptor entry(retry_eintr([&] { + return openat(root_fd, name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + 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 auto error = validate_tree(entry.get(), files); + if (error == Error::Ok) { + return EntryState::Valid; + } + if (error == Error::InvalidExternalData) { + return EntryState::Corrupt; + } + return error; +} + +struct StagingDirectory { + NSURL* url; + ~StagingDirectory() { + if (url != nil) { + NSError* error = nil; + if (![NSFileManager.defaultManager removeItemAtURL:url error:&error]) { + ET_LOG( + Error, + "Cannot remove Core AI staging directory: %s", + error.localizedDescription.UTF8String); + } + } + } +}; } // namespace Result parse_manifest(NSData* data) { @@ -293,4 +413,126 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return selected; } +Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* named_data, + NSString* staging_root, + NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + relative_path(manifest.path), + InvalidProgram, + "Core AI bundle must be selected before preparing its source"); + ET_CHECK_OR_RETURN_ERROR( + nonempty_string(key) && key.length == 64 && + [key rangeOfCharacterFromSet: + [[NSCharacterSet characterSetWithCharactersInString: + @"0123456789abcdef"] invertedSet]] + .location == NSNotFound, + InvalidArgument, + "Core AI staging key must be a 64-character lowercase hex digest"); + auto prepared = prepare_storage_root(staging_root, false); + if (!prepared.ok()) { + return prepared.error(); + } + NSURL* root = [NSURL fileURLWithPath:prepared.get() isDirectory:YES]; + ET_CHECK_OK_OR_RETURN_ERROR(validate_files( + manifest.files, [NSSet setWithObject:manifest.path.lastPathComponent])); + NSDictionary* files = manifest.files; + NSURL* model_url = [[root URLByAppendingPathComponent:key isDirectory:YES] + URLByAppendingPathComponent:manifest.path.lastPathComponent + isDirectory:YES]; + // Inspect, stage and publish relative to one root descriptor. + FileDescriptor root_fd(retry_eintr([&] { + return open(root.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + root_fd.get() >= 0, AccessFailed, "Cannot open Core AI storage root"); + auto state = inspect_entry(root_fd.get(), key, files); + if (!state.ok()) { + return state.error(); + } + if (*state == EntryState::Valid) { + return model_url; + } + // Recovery sources may still be in use by a previously acquired SDK model. + ET_CHECK_OR_RETURN_ERROR( + *state == EntryState::Missing, + InvalidExternalData, + "Cannot replace a damaged bookmark recovery source"); + ET_CHECK_OR_RETURN_ERROR( + named_data != nullptr, + InvalidProgram, + "Missing NamedDataMap for Core AI source recovery"); + NSString* staging_name = + [@".staging-" stringByAppendingString:NSUUID.UUID.UUIDString]; + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { + return mkdirat( + root_fd.get(), staging_name.fileSystemRepresentation, 0700); + }) == 0, + AccessFailed, + "Cannot create Core AI staging directory"); + StagingDirectory staging{[root URLByAppendingPathComponent:staging_name + isDirectory:YES]}; + FileDescriptor staging_fd(retry_eintr([&] { + return openat( + root_fd.get(), + staging_name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + staging_fd.get() >= 0, + AccessFailed, + "Cannot open Core AI staging directory"); + for (NSString* file in files) { + @autoreleasepool { + NSString* data_key = + [NSString stringWithFormat:@"coreai/%@/%@", manifest.hash, file]; + auto buffer = named_data->get_data(data_key.UTF8String); + ET_CHECK_OR_RETURN_ERROR( + buffer.ok(), + InvalidExternalData, + "Missing selected Core AI named data: %s", + data_key.UTF8String); + ET_CHECK_OR_RETURN_ERROR( + buffer->size() == 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.get(), file, buffer->data(), buffer->size())); + } + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_tree(staging_fd.get(), files)); + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::Rename) == 0, + AccessFailed, + "Core AI source publication interrupted"); + if (retry_eintr([&] { + return renameatx_np( + root_fd.get(), + staging_name.fileSystemRepresentation, + root_fd.get(), + key.fileSystemRepresentation, + RENAME_EXCL); + }) == 0) { + staging.url = nil; + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(root_fd.get())); + return model_url; + } + ET_CHECK_OR_RETURN_ERROR( + errno == EEXIST, AccessFailed, "Cannot publish Core AI stored bundle"); + // Another loader published first; use its source only if it is complete. + auto winner = inspect_entry(root_fd.get(), key, files); + if (!winner.ok()) { + return winner.error(); + } + ET_CHECK_OR_RETURN_ERROR( + *winner == EntryState::Valid, + InvalidExternalData, + "Concurrently published Core AI source is incomplete"); + return model_url; +} + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h index 8a09e5068ac..2ecb5740e1a 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -31,5 +31,12 @@ struct BookmarkDirectory { ::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 11c1a190e0d..d2c36e5e712 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 #include @@ -80,4 +81,79 @@ EXPECT_TRUE([NSFileManager.defaultManager removeItemAtURL:url error:nil]); } +::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; + } + 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..c637c817e40 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.h @@ -0,0 +1,53 @@ +/* + * 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" + +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; + 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 SetUp() override; + void check_concurrent_source(); + runtime::Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* data, + NSString* staging_root); + + private: + TestDirectory locks_; +}; + +} // 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..156c57a877a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm @@ -0,0 +1,94 @@ +/* + * 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 +#include +#include +#include +#include +#include "coreai_file.h" +#include "coreai_storage.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 Error::AccessFailed; + 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"; +} + +void CoreAISourceTest::SetUp() { + ASSERT_NE(locks_.url, nil); + struct stat info; + ASSERT_EQ(lstat(locks_.url.fileSystemRepresentation, &info), 0); + ASSERT_TRUE(S_ISDIR(info.st_mode)); +} + +Result CoreAISourceTest::prepare_source_bundle(const Manifest& manifest, + const runtime::NamedDataMap* data, + NSString* staging_root) { + NSString* key = + manifest.path == nil ? nil : manifest.bundle_digests[manifest.path.lastPathComponent]; + FileDescriptor lock(-1); + if (key != nil) { + FileDescriptor parent( + open(locks_.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC)); + if (parent.get() < 0) { + ADD_FAILURE() << "Cannot open source fixture lock root: " << strerror(errno); + return Error::AccessFailed; + } + NSString* name = [key stringByAppendingString:@".lock"]; + lock.reset(open_storage_shared_file(parent.get(), name.fileSystemRepresentation)); + if (lock.get() < 0 || retry_eintr([&] { return flock(lock.get(), LOCK_EX); }) != 0) { + ADD_FAILURE() << "Cannot acquire source fixture lock: " << strerror(errno); + return Error::AccessFailed; + } + } + return executorch::backends::coreai::prepare_source_bundle(manifest, data, staging_root, key); +} + +} // 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..58ea949e955 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_test.mm @@ -0,0 +1,385 @@ +/* + * 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 "coreai_bookmarks.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +namespace { +Result source_manifest(bool aot = false) { + auto parsed = parse_manifest(encode(aot ? aot_manifest_dict() : manifest_dict())); + if (!parsed.ok()) return parsed.error(); + return select_assets(parsed.get(), @"arch_b", @"macOS"); +} +} // namespace + +TEST_F(CoreAISourceTest, RejectsUnselectedAotSource) { + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + ASSERT_TRUE(parsed->aot_compiled); + ASSERT_EQ(parsed->path, nil); + EXPECT_FALSE(prepare_source_bundle(parsed.get(), nullptr, nil).ok()); +} + +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 Manifest* 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:storage.url.path error:nil] + .count, + 1u); + } +} + +TEST_F(CoreAISourceTest, ConcurrentColdAndWarmSource) { + check_concurrent_source(); +} + +TEST_F(CoreAISourceTest, RejectsInvalidKeysAndMissingSourceInputs) { + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSString* valid_key = fixture_identity('a'); + for (NSString* key in @[ @"../escape", [valid_key substringToIndex:63] ]) { + SCOPED_TRACE(key.UTF8String); + TestDirectory invalid_storage; + ASSERT_NE(invalid_storage.url, nil); + TestData invalid_data; + EXPECT_EQ(executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &invalid_data, invalid_storage.url.path, key) + .error(), + Error::InvalidArgument); + EXPECT_EQ(invalid_data.attempts, 0); + } + + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + EXPECT_FALSE(prepare_source_bundle(manifest.get(), nullptr, storage.url.path).ok()); + NSURL* absent = [storage.url URLByAppendingPathComponent:@"absent"]; + EXPECT_FALSE(prepare_source_bundle(manifest.get(), &data, absent.path).ok()); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:absent.path]); + auto recovered = prepare_source_bundle(manifest.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 = parse_manifest(encode(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 = parse_manifest(encode(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 = parse_manifest(encode(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 = parse_manifest(encode(dict)); + ASSERT_TRUE(manifest.ok()); + data.fail_at = 2; + EXPECT_FALSE(prepare_source_bundle(manifest.get(), &data, storage.url.path).ok()); + EXPECT_EQ( + [NSFileManager.defaultManager contentsOfDirectoryAtPath: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: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 = parse_manifest(encode(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) { + BookmarkDirectory directory; + ASSERT_NE(directory.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto root = resolve_bookmark_root(directory.url.path); + ASSERT_TRUE(root.ok()); + NSString* key = fixture_identity('9'); + auto lock = lock_bookmark(root.get(), key); + ASSERT_TRUE(lock.ok()); + auto staging = prepare_bookmark_staging(*lock.get()); + ASSERT_TRUE(staging.ok()); + auto first = executorch::backends::coreai::prepare_source_bundle(manifest.get(), &data, + staging.get(), key); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 1); + EXPECT_TRUE(([first.get().path + isEqualToString:[root.get() + stringByAppendingPathComponent: + [NSString stringWithFormat:@"staging/%@/model.aimodel", key]]])); + 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(manifest.get(), &data, + staging.get(), key); + 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, 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 = storage.url; + NSURL* final = + [root URLByAppendingPathComponent:manifest->bundle_digests[manifest->path.lastPathComponent]]; + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(StorageOperation::Rename); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + if (operation != StorageOperation::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->path.lastPathComponent]]); + EXPECT_EQ(publications, 1); + EXPECT_EQ(fault.hits, 0); + } + 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