diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index a5d6cf1d8fa..de6ad63b6cd 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -123,6 +123,11 @@ void XNNWeightsCache::reset_for_fresh_write() { if (is_heap_backed) { ++it; } else { + // Null the dangling slot before dropping the entry: the region it + // pointed into was just munmapped above. + if (it->second.offset < packed_data_ptrs_.size()) { + packed_data_ptrs_[it->second.offset] = nullptr; + } it = name_to_packed_data_metadata_.erase(it); } } @@ -738,6 +743,28 @@ bool XNNWeightsCache::load_packed_cache() { name_to_packed_data_metadata_[name] = meta; } + // The loop above stops as soon as fewer than 4 bytes remain, even if it + // has consumed fewer than entry_count entries. A truncated index that + // leaves <4 trailing bytes would otherwise be accepted as a partial cache, + // and the next save_packed_index would rewrite a trailer covering only that + // subset — permanently dropping the missing entries. Require that exactly + // entry_count entries were read and that the cursor consumed the whole + // index region; otherwise roll back like the in-loop corruption branches. + if (name_to_packed_data_metadata_.size() != entry_count || cursor != end) { + ET_LOG( + Error, + "load_packed_cache: index truncated (read %zu of %u entries, %zd trailing bytes); aborting load", + name_to_packed_data_metadata_.size(), + entry_count, + static_cast(end - cursor)); + munmap(map, file_size); + mmap_regions_.pop_back(); + name_to_packed_data_metadata_.clear(); + packed_data_ptrs_.clear(); + ptr_to_file_offset_.clear(); + return false; + } + // Start writing new packs AFTER the existing trailer (not at the // trailer's offset). Writing at index_start would overwrite the // valid trailer this PTE just loaded, breaking cross-process reuse diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index f4cc3fdad16..b399115031b 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -30,6 +30,11 @@ using executorch::runtime::FreeableBuffer; using executorch::runtime::MemoryAllocator; using executorch::runtime::Result; +// Grants the unit test access to private members so it can exercise the +// internal fresh-write / load paths directly. Test-only; no production code +// names this type. +class XNNWeightsCacheTestPeer; + struct PackedDataMeta { size_t offset{}; size_t data_size{0}; @@ -163,6 +168,8 @@ class XNNWeightsCache { } private: + friend class XNNWeightsCacheTestPeer; + static constexpr uint32_t kCacheMagic = 0x58505743; // "XPWC" // Bump when the on-disk layout (footer or per-entry record) changes. // v2: per-entry seed added — old v1 files don't carry seeds and would diff --git a/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp b/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp index d9661eae588..2ef7e8185ba 100644 --- a/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp +++ b/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp @@ -1064,4 +1064,127 @@ TEST_F(XNNWeightsCacheTest, ConcurrentOptionsAndSave_NoCrash_FileStable) { ::unlink(cache_path.c_str()); } +// Test seam: XNNWeightsCache befriends this type (see XNNWeightsCache.h) so +// the test can drive the private fresh-write path and inspect the +// offset->pointer table directly. +namespace executorch { +namespace backends { +namespace xnnpack { +namespace delegate { +class XNNWeightsCacheTestPeer { + public: + static void reset_for_fresh_write(XNNWeightsCache& c) { + c.reset_for_fresh_write(); + } + static void* offset_to_addr(XNNWeightsCache& c, size_t offset) { + return XNNWeightsCache::offset_to_addr(&c, offset); + } + static size_t offset_of(XNNWeightsCache& c, const std::string& name) { + return c.name_to_packed_data_metadata_.at(name).offset; + } +}; +} // namespace delegate +} // namespace xnnpack +} // namespace backends +} // namespace executorch + +using executorch::backends::xnnpack::delegate::XNNWeightsCacheTestPeer; + +TEST_F(XNNWeightsCacheTest, OffsetToAddr_AfterResetForFreshWrite_NotDangling) { + std::string cache_path = std::string("/tmp/xnn_weights_cache_dangling_") + + std::to_string(::getpid()) + ".packed_cache"; + ::unlink(cache_path.c_str()); + + std::vector batches{1, 2, 3}; + size_t input_channels = 3; + size_t output_channels = 4; + size_t num_batches = 1 * 2 * 3; + size_t padding = 32; + std::vector input(num_batches * input_channels + padding, 1.0f); + std::vector output(num_batches * output_channels, 0.0f); + + XNNWeightsCache cache; + cache.set_packed_cache_path(cache_path); + cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()); + BuildAndRunGraphWithWeightsCache( + cache, + batches, + input_channels, + output_channels, + input.data(), + output.data()); + ASSERT_EQ(cache.get_packed_data_names().size(), 1u); + + size_t offset = XNNWeightsCacheTestPeer::offset_of(cache, "weightbias"); + void* before = XNNWeightsCacheTestPeer::offset_to_addr(cache, offset); + ASSERT_NE(before, nullptr); + // The pointer is into a live MAP_SHARED region and is readable now. + ASSERT_NO_FATAL_FAILURE({ + volatile char c = *static_cast(before); + (void)c; + }); + + // Fresh-write reset: munmaps the region backing `before`. + XNNWeightsCacheTestPeer::reset_for_fresh_write(cache); + + // The slot must no longer reference the unmapped region. A nullptr return is + // safe (look_up_or_insert treats it as a miss); a non-null return would be a + // dangling pointer into munmapped memory. + void* after = XNNWeightsCacheTestPeer::offset_to_addr(cache, offset); + EXPECT_EQ(after, nullptr) + << "offset_to_addr returned a dangling pointer into a region that " + "reset_for_fresh_write() already munmapped"; + + ::unlink(cache_path.c_str()); +} + +TEST_F(XNNWeightsCacheTest, LoadPackedCache_RejectsMidTrailerTruncation) { + std::string cache_path = std::string("/tmp/xnn_weights_cache_midtrunc_") + + std::to_string(::getpid()) + ".packed_cache"; + ::unlink(cache_path.c_str()); + + // Layout (little-endian, matching XNNWeightsCache's on-disk format): + // [0, 16) 16 bytes of packed-data region (entry 0's bytes live here) + // [16, 41) entry 0: name_len(4)=1, name(1)="w", + // file_offset(8)=0, data_size(8)=8, seed(4)=0x12345678 + // [41, 43) 2 trailing bytes <- < 4 bytes, so the load loop bails here + // [43, 63) footer: index_start(8)=16, entry_count(4)=2, + // magic(4)="XPWC", version(4)=2 + // The footer is fully valid and claims 2 entries, but only 1 is present. + { + std::ofstream f(cache_path, std::ios::binary); + std::vector data_region(16, 0); + f.write(data_region.data(), data_region.size()); + + // entry 0 + write_le_u32(f, 1); // name_len + f.put('w'); // name + write_le_u64(f, 0); // file_offset + write_le_u64(f, 8); // data_size (<= index_start - file_offset) + write_le_u32(f, 0x12345678u); // seed + + // 2 trailing bytes: leaves < 4 bytes before the footer. + f.put('\0'); + f.put('\0'); + + // footer + write_le_u64(f, 16); // index_start + write_le_u32(f, 2); // entry_count (claims 2, only 1 present) + write_le_u32(f, 0x58505743u); // kCacheMagic "XPWC" + write_le_u32(f, 2); // kCacheVersion + } + + XNNWeightsCache cache; + cache.set_packed_cache_path(cache_path); + Error err = + cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()); + ASSERT_EQ(err, Error::Ok); + + EXPECT_EQ(cache.get_packed_data_names().size(), 0u) + << "mid-trailer-truncated cache (footer entry_count=2, only 1 entry " + "present) was wrongly accepted as a valid partial cache"; + + ::unlink(cache_path.c_str()); +} + #endif