poc/harness_flat_tensor_fuzzer.cpp

// libFuzzer harness for executorch::extension::FlatTensorDataMap::load()
// Targets the .ptd flatbuffer parsing path (extension/flat_tensor/flat_tensor_data_map.cpp).
// Authorized local testing only - huntr MFV scope: ExecuTorch .pte/.ptd parser.

#include <cstddef>
#include <cstdint>

#include <executorch/extension/data_loader/buffer_data_loader.h>
#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
#include <executorch/runtime/platform/runtime.h>

using executorch::extension::BufferDataLoader;
using executorch::extension::FlatTensorDataMap;

static bool g_initialized = false;

extern "C" int LLVMFuzzerTestOneInput(const std::uint8_t* data, std::size_t size) {
  if (!g_initialized) {
    executorch::runtime::runtime_init();
    g_initialized = true;
  }

  constexpr std::size_t kMaxInput = 32U * 1024U * 1024U;
  if (data == nullptr || size == 0 || size > kMaxInput) {
    return 0;
  }

  BufferDataLoader loader(data, size);
  auto map = FlatTensorDataMap::load(&loader);
  if (!map.ok()) {
    return 0;
  }

  auto& m = map.get();
  auto num_keys_res = m.get_num_keys();
  if (num_keys_res.ok()) {
    uint32_t n = num_keys_res.get();
    // Bound the loop - malicious n could be huge but get_key() range-checks it.
    uint32_t iter = n > 4096 ? 4096 : n;
    for (uint32_t i = 0; i < iter; ++i) {
      auto key_res = m.get_key(i);
      if (key_res.ok()) {
        auto layout = m.get_tensor_layout(executorch::aten::string_view(key_res.get()));
        if (layout.ok()) {
          (void)layout.get().nbytes();
        }
        auto data_res = m.get_data(executorch::aten::string_view(key_res.get()));
        if (data_res.ok()) {
          data_res.get().Free();
        }
      }
    }
  }

  return 0;
}