// Copyright 2017 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "sandbox/win/src/heap_helper.h"

#include <windows.h>

#include "base/containers/heap_array.h"
#include "base/logging.h"
#include "base/memory/raw_ptr_exclusion.h"
#include "base/memory/ref_counted.h"

namespace sandbox {
namespace {
#pragma pack(1)

// These are undocumented, but readily found on the internet.
constexpr DWORD kHeapClass8 = 0x00008000;  // CSR port heap
constexpr DWORD kHeapClassMask = 0x0000f000;

constexpr DWORD kHeapSegmentSignature = 0xffeeffee;
constexpr DWORD kHeapSignature = 0xeeffeeff;

typedef struct _HEAP_ENTRY {
  PVOID Data1;
  PVOID Data2;
} HEAP_ENTRY, *PHEAP_ENTRY;

// The _HEAP struct is not documented, so char arrays are used to space out the
// struct of the fields that are not relevant. However, this spacing is
// different because of the different pointer widths between 32 and 64-bit.
// So 32 and 64 bit structs are defined.
struct _HEAP_32 {
  HEAP_ENTRY HeapEntry;
  DWORD SegmentSignature;
  DWORD SegmentFlags;
  LIST_ENTRY SegmentListEntry;
  // `Heap` is not a raw_ptr<...>, because reinterpret_cast of uninitialized
  // memory to raw_ptr can cause ref-counting mismatch.
  RAW_PTR_EXCLUSION struct _HEAP_32* Heap;
  char Unknown0[0x24];
  // Offset 0x40
  DWORD Flags;
  // Offset 0x60
  char Unknown1[0x1c];
  DWORD Signature;
  // Other stuff that is not relevant.
};

struct _HEAP_64 {
  HEAP_ENTRY HeapEntry;
  DWORD SegmentSignature;
  DWORD SegmentFlags;
  LIST_ENTRY SegmentListEntry;
  // `Heap` is not a raw_ptr<...>, because reinterpret_cast of uninitialized
  // memory to raw_ptr can cause ref-counting mismatch.
  RAW_PTR_EXCLUSION struct _HEAP_64* Heap;
  char Unknown0[0x40];
  // Offset 0x70
  DWORD Flags;
  // Offset 0x98
  char Unknown1[0x24];
  DWORD Signature;
  // Other stuff that is not relevant.
};

#if defined(_WIN64)
using _HEAP = _HEAP_64;
#else   // defined(_WIN64)
using _HEAP = _HEAP_32;
#endif  // defined(_WIN64)

bool ValidateHeap(_HEAP* heap) {
  if (heap->SegmentSignature != kHeapSegmentSignature)
    return false;
  if (heap->Heap != heap)
    return false;
  if (heap->Signature != kHeapSignature)
    return false;
  return true;
}

}  // namespace

bool HeapFlags(HANDLE handle, DWORD* flags) {
  if (!handle || !flags) {
    // This is an error.
    return false;
  }
  _HEAP* heap = reinterpret_cast<_HEAP*>(handle);
  if (!ValidateHeap(heap)) {
    DLOG(ERROR) << "unable to validate heap";
    return false;
  }
  *flags = heap->Flags;
  return true;
}

HANDLE FindCsrPortHeap() {
  DWORD number_of_heaps = ::GetProcessHeaps(0, nullptr);
  auto all_heaps = base::HeapArray<HANDLE>::Uninit(number_of_heaps);
  if (::GetProcessHeaps(number_of_heaps, all_heaps.data()) != number_of_heaps) {
    return nullptr;
  }

  // Search for the CSR port heap handle, identified purely based on flags.
  HANDLE csr_port_heap = nullptr;
  for (HANDLE handle : all_heaps) {
    DWORD flags = 0;
    if (!HeapFlags(handle, &flags)) {
      DLOG(ERROR) << "Unable to get flags for this heap";
      continue;
    }
    if ((flags & kHeapClassMask) == kHeapClass8) {
      if (csr_port_heap) {
        DLOG(ERROR) << "Found multiple suitable CSR Port heaps";
        return nullptr;
      }
      csr_port_heap = handle;
    }
  }
  return csr_port_heap;
}

}  // namespace sandbox