/* radare2 - LGPL - Copyright 2025 - pancake */

#include <r_fs.h>
#include <r_types.h>

#define GPT_SIGNATURE "EFI PART"
#define GPT_SIGNATURE_LEN 8

R_PACKED(
	typedef struct {
		ut8 signature[8];
		ut32 revision;
		ut32 header_size;
		ut32 crc32;
		ut32 reserved;
		ut64 current_lba;
		ut64 backup_lba;
		ut64 first_usable_lba;
		ut64 last_usable_lba;
		ut8 disk_guid[16];
		ut64 partition_entries_lba;
		ut32 num_partition_entries;
		ut32 partition_entry_size;
		ut32 partition_entries_crc32;
	})
GPTHeader;

R_PACKED(
	typedef struct {
		ut8 type_guid[16];
		ut8 unique_guid[16];
		ut64 first_lba;
		ut64 last_lba;
		ut64 attributes;
		ut16 name[36]; // UTF-16
	})
GPTEntry;

static int fs_part_gpt(void *disk, void *ptr, void *closure) {
	GPTHeader header;
	RFS *fs = (RFS *)disk;
	RFSPartition *par = NULL;
	RFSPartitionIterator iterate = (RFSPartitionIterator)ptr;
	RList *list = (RList *)closure;

	// Read GPT header at LBA 1
	if (fs->iob.read_at (fs->iob.io, 512, (ut8 *)&header, sizeof (header)) != sizeof (header)) {
		R_LOG_ERROR ("Failed to read GPT header");
		return 0;
	}
	if (memcmp (header.signature, GPT_SIGNATURE, GPT_SIGNATURE_LEN) != 0) {
		R_LOG_ERROR ("Invalid GPT signature");
		return 0;
	}

	const ut8 *header_bytes = (const ut8 *)&header;
	ut64 entries_lba = r_read_le64 (header_bytes + 72) * 512;
	ut32 num_entries = r_read_le32 (header_bytes + 80);
	ut32 entry_size = r_read_le32 (header_bytes + 84);

	if (entry_size != 128) {
		R_LOG_ERROR ("Unsupported partition entry size: %u", entry_size);
		return 0;
	}

	// Bounds checking for num_entries and multiplication overflow
	if (num_entries == 0 || num_entries > 1024) {
		R_LOG_ERROR ("Invalid number of partition entries: %u", num_entries);
		return 0;
	}
	if (entry_size == 0 || entry_size > 4096) {
		R_LOG_ERROR ("Invalid partition entry size: %u", entry_size);
		return 0;
	}
	if ((size_t)entry_size > SIZE_MAX / (size_t)num_entries) {
		R_LOG_ERROR ("Partition entries allocation size overflow");
		return 0;
	}
	size_t alloc_size = (size_t)num_entries *(size_t)entry_size;
	if (alloc_size > (4 * 1024 * 1024)) { // 4MB sanity check
		R_LOG_ERROR ("Partition entries allocation size too large: %zu", alloc_size);
		return 0;
	}
	ut8 *entries = malloc (alloc_size);
	if (!entries) {
		R_LOG_ERROR ("Failed to allocate memory for GPT partition entries");
		return 0;
	}
	if (fs->iob.read_at (fs->iob.io, entries_lba, entries, (int)alloc_size) != (int)alloc_size) {
		R_LOG_ERROR ("Failed to read GPT partition entries");
		free (entries);
		return 0;
	}

	int i;
	for (i = 0; i < num_entries; i++) {
		const ut8 *e = &entries[i * entry_size];
		// Check if type_guid is not all zeros
		bool is_empty = true;
		int j;
		for (j = 0; j < 16; j++) {
			if (e[j] != 0) {
				is_empty = false;
				break;
			}
		}
		if (!is_empty) {
			ut64 first_lba = r_read_le64 (e + 32);
			ut64 last_lba = r_read_le64 (e + 40);
			ut64 start = first_lba * 512;
			ut64 length = (last_lba - first_lba + 1) * 512;
			par = r_fs_partition_new (i, start, length);
			par->type = 0; // GPT doesn't have byte type, maybe use index or something
			iterate (disk, par, list);
		}
	}

	free (entries);
	return 0;
}