//go:build windows

package winpwn

import (
	"fmt"
	"unsafe"

	"golang.org/x/sys/windows"
)

// Offsets within ntdll's x64 _PEB, confirmed via `dt ntdll!_PEB` the same
// way heap.go's _HEAP offsets were -- see heap.go's top comment.
const (
	pebOffNumberOfHeaps = 0x0e8
	pebOffProcessHeaps  = 0x0f0 // PVOID*, an array of NumberOfHeaps heap addresses
)

// ListProcessHeaps enumerates every heap that exists in pid's address
// space by walking PEB.ProcessHeaps from outside the process -- the exact
// same array GetProcessHeaps() reads from inside one, just reached via
// ReadProcessMemory the way ResolveModuleBase (procmem_windows.go) reads
// PEB.Ldr for the loaded-module list instead of needing a leak. The
// default process heap (PEB.ProcessHeap) is always included, since
// HeapCreate registers every heap -- including the default one ntdll
// creates before main() even runs -- into this same array.
//
// This is the natural companion to DetectHeapKind/ReadHeap: once you have
// a PID and nothing else, ListProcessHeaps is how you find an address
// worth handing to either of them, instead of needing a leaked heap handle
// from the target's own output first.
func ListProcessHeaps(pid uint32) ([]uint64, error) {
	mem, err := OpenProcessMemory(pid, 0)
	if err != nil {
		return nil, err
	}
	defer mem.Close()

	var pbi windows.PROCESS_BASIC_INFORMATION
	var retLen uint32
	if err := windows.NtQueryInformationProcess(mem.Handle, windows.ProcessBasicInformation,
		unsafe.Pointer(&pbi), uint32(unsafe.Sizeof(pbi)), &retLen); err != nil {
		return nil, fmt.Errorf("NtQueryInformationProcess(ProcessBasicInformation): %w", err)
	}
	pebAddr := uint64(uintptr(unsafe.Pointer(pbi.PebBaseAddress)))
	if pebAddr == 0 {
		return nil, fmt.Errorf("PEB address for pid %d is null", pid)
	}

	numHeaps, err := readUint32AtValue(mem, int64(pebAddr)+pebOffNumberOfHeaps)
	if err != nil {
		return nil, fmt.Errorf("reading PEB.NumberOfHeaps: %w", err)
	}
	arrayAddr, err := readUint64At(mem, int64(pebAddr)+pebOffProcessHeaps)
	if err != nil {
		return nil, fmt.Errorf("reading PEB.ProcessHeaps: %w", err)
	}

	heaps := make([]uint64, 0, numHeaps)
	for i := uint32(0); i < numHeaps; i++ {
		addr, err := readUint64At(mem, int64(arrayAddr)+int64(i)*8)
		if err != nil {
			return heaps, fmt.Errorf("reading ProcessHeaps[%d] (of %d): %w", i, numHeaps, err)
		}
		heaps = append(heaps, addr)
	}
	return heaps, nil
}