/**
 * PoC: MemoryReadAdapter heap out-of-bounds read
 *
 * This test directly demonstrates that caffe2::serialize::MemoryReadAdapter::read()
 * performs NO bounds checking. The read() method blindly memcpy's from data_+pos
 * for n bytes, without checking that pos+n <= size_.
 *
 * Compile with ASAN:
 *   g++ -fsanitize=address -g -I$(python3 -c "import torch; print(torch.utils.cmake_prefix_path)")/../../include \
 *       test_oob_read.cpp -o test_oob_read
 *
 * Expected output: ASAN reports heap-buffer-overflow
 */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>

// Inline the vulnerable class (exact copy from caffe2/serialize/in_memory_adapter.h)
// to avoid needing the full PyTorch build
class MemoryReadAdapter {
 public:
  explicit MemoryReadAdapter(const void* data, int64_t size)
      : data_(data), size_(size) {}

  size_t size() const {
    return size_;
  }

  // THIS IS THE VULNERABILITY: no bounds checking on pos or n vs size_
  size_t read(uint64_t pos, void* buf, size_t n, const char* what = "") const {
    (void)what;
    memcpy(buf, (int8_t*)(data_) + pos, n);  // NO CHECK: pos+n vs size_
    return n;
  }

 private:
  const void* data_;
  int64_t size_;
};

// For comparison: miniz's own memory reader HAS bounds checking
// (from third_party/miniz-3.0.2/miniz.c, mz_zip_mem_read_func)
size_t safe_read(const void* data, size_t data_size, uint64_t pos, void* buf, size_t n) {
    size_t s = (pos >= data_size) ? 0 : (size_t)((data_size - pos < n) ? data_size - pos : n);
    memcpy(buf, (const uint8_t*)data + pos, s);
    return s;
}

int main() {
    // Allocate a small buffer (32 bytes)
    const size_t BUF_SIZE = 32;
    char* data = (char*)malloc(BUF_SIZE);
    if (!data) return 1;
    memset(data, 'A', BUF_SIZE);

    // Create MemoryReadAdapter with correct size
    MemoryReadAdapter adapter(data, BUF_SIZE);
    printf("Buffer size: %zu bytes at %p\n", adapter.size(), data);

    char output[256];
    memset(output, 0, sizeof(output));

    // Test 1: Normal read within bounds (should work fine)
    printf("\n[Test 1] Reading 16 bytes at offset 0 (within bounds)...\n");
    adapter.read(0, output, 16);
    printf("  OK: read 16 bytes\n");

    // Test 2: Read past the end of the buffer (OOB!)
    printf("\n[Test 2] Reading 64 bytes at offset 0 (32 bytes past buffer end)...\n");
    printf("  Buffer is %zu bytes, but requesting 64 bytes\n", BUF_SIZE);
    printf("  MemoryReadAdapter::read() will memcpy 64 bytes - reading 32 bytes of HEAP DATA\n");
    adapter.read(0, output, 64);  // ASAN: heap-buffer-overflow
    printf("  Leaked %zu bytes past buffer end!\n", (size_t)64 - BUF_SIZE);

    // Test 3: Read at offset past buffer end
    printf("\n[Test 3] Reading 16 bytes at offset 128 (entirely past buffer)...\n");
    adapter.read(128, output, 16);  // ASAN: heap-buffer-overflow
    printf("  Read from offset 128, buffer is only %zu bytes!\n", BUF_SIZE);

    free(data);
    return 0;
}