Fix use-after-free of unpacked constants in the XNNPACK weights cache

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -2234,6 +2234,10 @@ ET_NODISCARD Error XNNCompiler::compileModel(
   Error err = Error::Ok;
   for (auto value : *flatbuffer_graph->xvalues()) {
     size_t prev_buffers = unpacked_buffers.size();
+    // With the weights cache the buffers land in the cache rather than in
+    // unpacked_buffers, so track its list too.
+    size_t prev_cached_buffers =
+        use_weight_cache ? weights_cache->get_num_unpacked_data() : 0;
     err = defineTensor(
         subgraph.get(),
         remapped_ids,
@@ -2262,6 +2266,10 @@ ET_NODISCARD Error XNNCompiler::compileModel(
         executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i]));
       }
       unpacked_buffers.resize(prev_buffers);
+      if (use_weight_cache) {
+        weights_cache->take_unpacked_data_from(
+            prev_cached_buffers, executor->unpacked_buffers_);
+      }
     }
   }
 
@@ -2310,6 +2318,8 @@ ET_NODISCARD Error XNNCompiler::compileModel(
 
   std::vector<std::string> packed_weights_names;
   if (use_weight_cache) {
+    // Constants XNNPACK does not pack were already moved to the executor in
+    // the value loop above, so everything left here is safe to free.
     auto packed_weights_names_result = weights_cache->finalize_for_runtime();
     ET_CHECK_OR_RETURN_ERROR(
         packed_weights_names_result.ok(),

backends/xnnpack/runtime/XNNWeightsCache.cpp

@@ -215,6 +215,18 @@ Error XNNWeightsCache::initialize_for_runtime(
   named_data_map_ = named_data_map;
   is_finalized_ = false;
 
+  // An earlier compile can bail out between its first load_unpacked_data and
+  // its finalize_for_runtime, leaving its constants here. A successful compile
+  // frees whatever is left in finalize_for_runtime, but a failed one never
+  // reaches it, so back to back failures would pile up on an instance that
+  // lives as long as the process. The model that loaded them is gone, so free
+  // them here.
+  for (FreeableBuffer& buffer : unpacked_data_) {
+    buffer.Free();
+  }
+  unpacked_data_.clear();
+  unpacked_data_to_name_.clear();
+
 #ifndef _WIN32
   if (packed_cache_path_.empty() || packed_file_fd_ >= 0) {
     return Error::Ok;
@@ -296,11 +308,35 @@ Error XNNWeightsCache::initialize_for_runtime(
   return Error::Ok;
 }
 
+void XNNWeightsCache::take_unpacked_data_from(
+    size_t first_index,
+    std::vector<FreeableBuffer>& out) {
+  // No bounds check: the loop below is empty when first_index is at or past
+  // the end, and the list only grows during a compile so that cannot happen
+  // anyway. An index that is too small cannot be detected here, and would hand
+  // the executor buffers of values XNNPACK did pack; callers must pass the
+  // get_num_unpacked_data() taken immediately before the value was defined.
+  for (size_t i = first_index; i < unpacked_data_.size(); i++) {
+    // The name map is keyed on the data pointer, which a move preserves, so
+    // look_up_or_insert keeps naming packed entries correctly after this.
+    out.push_back(std::move(unpacked_data_[i]));
+  }
+  // The moved-from entries stay in the list rather than being erased:
+  // FreeableBuffer deletes move assignment, so vector::erase does not compile,
+  // and a moved-from buffer holds a null pointer, which makes the Free() in
+  // finalize_for_runtime a no-op. Leaving them also keeps the indices handed
+  // out by get_num_unpacked_data() stable across calls.
+}
+
 Result<std::vector<std::string>> XNNWeightsCache::finalize_for_runtime() {
   is_finalized_ = true;
 
-  // All data has been packed by create_runtime
-  // so we clear the unpacked data as it is no longer needed
+  // Everything still here belongs to a value XNNPACK packed, so create_runtime
+  // has copied it into the packed region and the unpacked copy can go. Buffers
+  // for values XNNPACK does not pack (PReLU slopes, for example) were handed to
+  // the executor by take_unpacked_data_from while the graph was being built;
+  // the subgraph keeps pointers into those, so freeing them here would leave
+  // the runtime reading freed memory.
   for (FreeableBuffer& buffer : unpacked_data_) {
     buffer.Free();
   }

backends/xnnpack/runtime/XNNWeightsCache.h

@@ -91,6 +91,26 @@ class XNNWeightsCache {
    */
   Result<std::vector<std::string>> finalize_for_runtime();
 
+  /**
+   * Transfers ownership of the unpacked buffers loaded since `first_index`
+   * out of this cache. finalize_for_runtime() will not free them.
+   *
+   * For values XNNPACK does not pack (PReLU slopes, for example) the subgraph
+   * keeps a pointer into the unpacked memory for the life of the runtime, so
+   * something has to keep those buffers alive past finalize_for_runtime().
+   * Callers pair this with get_num_unpacked_data() taken before the value was
+   * defined, which is how the non-weights-cache path in XNNCompiler decides
+   * what to retain.
+   *
+   * @param[in] first_index Index into the unpacked buffer list, from
+   *     get_num_unpacked_data() before the value was defined.
+   * @param[out] out Receives the buffers. The caller owns them and must keep
+   *     them alive for at least as long as the runtime.
+   */
+  void take_unpacked_data_from(
+      size_t first_index,
+      std::vector<FreeableBuffer>& out);
+
   // Taken from XNN_ALLOCATION_ALIGNMENT in xnnpack/common.h
   static const size_t kPackedAllocationAlignment = 64;
 

backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp

@@ -320,6 +320,54 @@ TEST_F(XNNWeightsCacheTest, ReusePackedWeights) {
   ASSERT_EQ(packed_data_names.size(), 0);
 }
 
+TEST_F(XNNWeightsCacheTest, TakenUnpackedDataSurvivesFinalize) {
+  // XNNPACK does not pack every constant it is given: PReLU slopes, for one,
+  // stay as the unpacked buffer the subgraph points at. Those buffers are taken
+  // out of the cache while the graph is built, and finalize_for_runtime must
+  // leave them alone while still freeing the ones it packed.
+  XNNWeightsCache cache;
+  cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get());
+
+  Result<const uint8_t*> packed = cache.load_unpacked_data("weight");
+  ASSERT_EQ(packed.error(), Error::Ok);
+  // The index the caller would take immediately before defining the value it
+  // wants to keep, which is what XNNCompiler passes down.
+  size_t first_index = cache.get_num_unpacked_data();
+  Result<const uint8_t*> retained_load = cache.load_unpacked_data("bias");
+  ASSERT_EQ(retained_load.error(), Error::Ok);
+  ASSERT_EQ(cache.get_num_unpacked_data(), first_index + 1);
+
+  std::vector<executorch::runtime::FreeableBuffer> retained;
+  cache.take_unpacked_data_from(first_index, retained);
+  ASSERT_EQ(retained.size(), 1u);
+  ASSERT_EQ(retained[0].size(), static_cast<size_t>(kSegmentSizes[1]));
+
+  Result<std::vector<std::string>> names = cache.finalize_for_runtime();
+  ASSERT_EQ(names.error(), Error::Ok);
+
+  // Still readable, and still the "bias" segment, which SetUp fills with 2s.
+  ASSERT_NE(retained[0].data(), nullptr);
+  const uint8_t* data = static_cast<const uint8_t*>(retained[0].data());
+  for (size_t i = 0; i < retained[0].size(); i++) {
+    ASSERT_EQ(data[i], 2) << "byte " << i << " was freed or overwritten";
+  }
+}
+
+TEST_F(XNNWeightsCacheTest, InitializeForRuntimeClearsLeftoverUnpackedData) {
+  // A compile that bails out after loading a constant never reaches
+  // finalize_for_runtime. This instance outlives any one model, so the next
+  // initialize has to drop what the failed compile left behind.
+  XNNWeightsCache cache;
+  cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get());
+  Result<const uint8_t*> loaded = cache.load_unpacked_data("weight");
+  ASSERT_EQ(loaded.error(), Error::Ok);
+  ASSERT_EQ(cache.get_num_unpacked_data(), 1u);
+
+  // No finalize_for_runtime(): this is the failed-compile path.
+  cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get());
+  ASSERT_EQ(cache.get_num_unpacked_data(), 0u);
+}
+
 #ifndef _WIN32
 // Verify pack-and-run works when packed weight allocations go to a
 // MAP_SHARED file instead of heap. The cache path is unique per test so