/*
 * Copyright 2007-2013, Axel Dörfler, axeld@pinc-software.de.
 * Copyright 2009, Michael Lotz, mmlr@mlotz.ch. All rights reserved.
 *
 * Distributed under the terms of the MIT License.
 */


#include "Header.h"

#include <unistd.h>
#include <stdio.h>
#include <string.h>

#include <KernelExport.h>

#ifdef _KERNEL_MODE
#	include <util/kernel_cpp.h>
#else
#	include <new>
#endif

#ifndef _BOOT_MODE
#include "PartitionMap.h"
#include "PartitionMapWriter.h"
#endif

#if !defined(_BOOT_MODE) && !defined(_USER_MODE)
#include "uuid.h"
#endif

#include "crc32.h"
#include "utility.h"


#define TRACE_EFI_GPT
#ifdef TRACE_EFI_GPT
#	ifndef _KERNEL_MODE
#		define dprintf printf
#	endif
#	define TRACE(x) dprintf x
#else
#	define TRACE(x) ;
#endif


namespace EFI {


Header::Header(int fd, uint64 lastBlock, uint32 blockSize)
	:
	fBlockSize(blockSize),
	fStatus(B_NO_INIT),
	fEntries(NULL),
	fDirty(false)
{
	// TODO: check the correctness of the protective MBR and warn if invalid

	// Read and check the partition table header

	fStatus = _Read(fd, (uint64)EFI_HEADER_LOCATION * blockSize,
		&fHeader, sizeof(gpt_table_header));
	if (fStatus == B_OK) {
		if (!_IsHeaderValid(fHeader, EFI_HEADER_LOCATION))
			fStatus = B_BAD_DATA;
	}

	if (fStatus == B_OK && lastBlock != fHeader.AlternateBlock()) {
		dprintf("gpt: alternate header not in last block (%" B_PRIu64 " vs. %"
			B_PRIu64 ")\n", fHeader.AlternateBlock(), lastBlock);
		lastBlock = fHeader.AlternateBlock();
	}

	// Read backup header, too
	status_t status = _Read(fd, lastBlock * blockSize, &fBackupHeader,
		sizeof(gpt_table_header));
	if (status == B_OK) {
		if (!_IsHeaderValid(fBackupHeader, lastBlock))
			status = B_BAD_DATA;
	}

	// If both headers are invalid, bail out -- this is probably not a GPT disk
	if (status != B_OK && fStatus != B_OK)
		return;

	if (fStatus != B_OK) {
		// Recreate primary header from the backup
		fHeader = fBackupHeader;
		fHeader.SetAbsoluteBlock(EFI_HEADER_LOCATION);
		fHeader.SetEntriesBlock(EFI_PARTITION_ENTRIES_BLOCK);
		fHeader.SetAlternateBlock(lastBlock);
		fDirty = true;
	} else if (status != B_OK) {
		// Recreate backup header from primary
		_SetBackupHeaderFromPrimary(lastBlock);
	}

	// allocate, read, and check partition entry array

	fEntries = new (std::nothrow) uint8[_EntryArraySize()];
	if (fEntries == NULL) {
		// TODO: if there cannot be allocated enough (ie. the boot loader's
		//	heap is limited), try a smaller size before failing
		fStatus = B_NO_MEMORY;
		return;
	}

	fStatus = _Read(fd, fHeader.EntriesBlock() * blockSize,
		fEntries, _EntryArraySize());
	if (fStatus != B_OK || !_ValidateEntriesCRC()) {
		// Read backup entries instead
		fStatus = _Read(fd, fBackupHeader.EntriesBlock() * blockSize,
			fEntries, _EntryArraySize());
		if (fStatus != B_OK)
			return;

		if (!_ValidateEntriesCRC()) {
			fStatus = B_BAD_DATA;
			return;
		}
	}

	// TODO: check overlapping or out of range partitions

#ifdef TRACE_EFI_GPT
	_Dump(fHeader);
	_Dump(fBackupHeader);
	_DumpPartitions();
#endif

	fStatus = B_OK;
}


#if !defined(_BOOT_MODE) && !defined(_USER_MODE)
Header::Header(uint64 lastBlock, uint32 blockSize)
	:
	fBlockSize(blockSize),
	fStatus(B_NO_INIT),
	fEntries(NULL),
	fDirty(true)
{
	TRACE(("EFI::Header: Initialize GPT, block size %" B_PRIu32 "\n",
		blockSize));

	// Initialize to an empty header
	memcpy(fHeader.header, EFI_PARTITION_HEADER, sizeof(fHeader.header));
	fHeader.SetRevision(EFI_TABLE_REVISION);
	fHeader.SetHeaderSize(sizeof(fHeader));
	fHeader.SetHeaderCRC(0);
	fHeader.SetAbsoluteBlock(EFI_HEADER_LOCATION);
	fHeader.SetAlternateBlock(lastBlock);
	uuid_t uuid;
	uuid_generate_random(uuid);
	memcpy((uint8*)&fHeader.disk_guid, uuid, sizeof(guid_t));
	fHeader.SetEntriesBlock(EFI_PARTITION_ENTRIES_BLOCK);
	fHeader.SetEntryCount(EFI_PARTITION_ENTRY_COUNT);
	fHeader.SetEntrySize(EFI_PARTITION_ENTRY_SIZE);
	fHeader.SetEntriesCRC(0);

	size_t arraySize = _EntryArraySize();
	fEntries = new (std::nothrow) uint8[arraySize];
	if (fEntries == NULL) {
		fStatus = B_NO_MEMORY;
		return;
	}

	memset(fEntries, 0, arraySize);
		// TODO: initialize the entry guids

	uint32 entryBlocks = (arraySize + fBlockSize - 1) / fBlockSize;
	fHeader.SetFirstUsableBlock(EFI_PARTITION_ENTRIES_BLOCK + entryBlocks);
	fHeader.SetLastUsableBlock(lastBlock - 1 - entryBlocks);

	_SetBackupHeaderFromPrimary(lastBlock);

#ifdef TRACE_EFI_GPT
	_Dump(fHeader);
	_DumpPartitions();
#endif

	fStatus = B_OK;
}
#endif // !_BOOT_MODE && !_USER_MODE


Header::~Header()
{
	delete[] fEntries;
}


status_t
Header::InitCheck() const
{
	return fStatus;
}


bool
Header::IsDirty() const
{
	return fDirty;
}


#ifndef _BOOT_MODE
status_t
Header::WriteEntry(int fd, uint32 entryIndex)
{
	off_t entryOffset = entryIndex * fHeader.EntrySize();

	status_t status = _Write(fd,
		fHeader.EntriesBlock() * fBlockSize + entryOffset,
		fEntries + entryOffset, fHeader.EntrySize());
	if (status != B_OK)
		return status;

	// Update header, too -- the entries CRC changed
	status = _WriteHeader(fd);

	// Write backup
	status_t backupStatus = _Write(fd,
		fBackupHeader.EntriesBlock() * fBlockSize + entryOffset,
		fEntries + entryOffset, fHeader.EntrySize());

	if (status == B_OK && backupStatus == B_OK)
		fDirty = false;
	return status == B_OK ? backupStatus : status;
}


status_t
Header::Write(int fd)
{
	// Try to write the protective MBR
	PartitionMap partitionMap;
	PrimaryPartition *partition = NULL;
	uint32 index = 0;
	while ((partition = partitionMap.PrimaryPartitionAt(index)) != NULL) {
		if (index == 0) {
			uint64 deviceSize = fHeader.AlternateBlock() * fBlockSize;
			partition->SetTo(fBlockSize, deviceSize, 0xEE, false, fBlockSize);
		} else
			partition->Unset();
		++index;
	}
	PartitionMapWriter writer(fd, fBlockSize);
	writer.WriteMBR(&partitionMap, true);
		// We also write the bootcode, so we can boot GPT disks from BIOS

	status_t status = _Write(fd, fHeader.EntriesBlock() * fBlockSize, fEntries,
		_EntryArraySize());
	if (status != B_OK)
		return status;

	// First write the header, so that we have at least one completely correct
	// data set
	status = _WriteHeader(fd);


	// Write backup entries
	status_t backupStatus = _Write(fd,
		fBackupHeader.EntriesBlock() * fBlockSize, fEntries, _EntryArraySize());

	if (status == B_OK && backupStatus == B_OK)
		fDirty = false;
	return status == B_OK ? backupStatus : status;
}


status_t
Header::_WriteHeader(int fd)
{
	_UpdateCRC();

	status_t status = _Write(fd, fHeader.AbsoluteBlock() * fBlockSize,
		&fHeader, sizeof(gpt_table_header));
	if (status != B_OK)
		return status;

	return _Write(fd, fBackupHeader.AbsoluteBlock() * fBlockSize,
		&fBackupHeader, sizeof(gpt_table_header));
}


status_t
Header::_Write(int fd, off_t offset, const void* data, size_t size) const
{
	ssize_t bytesWritten = write_pos(fd, offset, data, size);
	if (bytesWritten < 0)
		return bytesWritten;
	if (bytesWritten != (ssize_t)size)
		return B_IO_ERROR;

	return B_OK;
}


void
Header::_UpdateCRC()
{
	_UpdateCRC(fHeader);
	_UpdateCRC(fBackupHeader);
}


void
Header::_UpdateCRC(gpt_table_header& header)
{
	header.SetEntriesCRC(crc32(fEntries, _EntryArraySize()));
	header.SetHeaderCRC(0);
	header.SetHeaderCRC(crc32((uint8*)&header, sizeof(gpt_table_header)));
}
#endif // !_BOOT_MODE


status_t
Header::_Read(int fd, off_t offset, void* data, size_t size) const
{
	ssize_t bytesRead = read_pos(fd, offset, data, size);
	if (bytesRead < 0)
		return bytesRead;
	if (bytesRead != (ssize_t)size)
		return B_IO_ERROR;

	return B_OK;
}


bool
Header::_IsHeaderValid(gpt_table_header& header, uint64 block)
{
	return !memcmp(header.header, EFI_PARTITION_HEADER, sizeof(header.header))
		&& _ValidateHeaderCRC(header)
		&& header.AbsoluteBlock() == block;
}


bool
Header::_ValidateHeaderCRC(gpt_table_header& header)
{
	uint32 originalCRC = header.HeaderCRC();
	header.SetHeaderCRC(0);

	bool matches = originalCRC == crc32((const uint8*)&header,
		sizeof(gpt_table_header));

	header.SetHeaderCRC(originalCRC);
	return matches;
}


bool
Header::_ValidateEntriesCRC() const
{
	return fHeader.EntriesCRC() == crc32(fEntries, _EntryArraySize());
}


void
Header::_SetBackupHeaderFromPrimary(uint64 lastBlock)
{
	fBackupHeader = fHeader;
	fBackupHeader.SetAbsoluteBlock(lastBlock);
	fBackupHeader.SetEntriesBlock(
		lastBlock - _EntryArraySize() / fBlockSize);
	fBackupHeader.SetAlternateBlock(1);
}


#ifdef TRACE_EFI_GPT
const char *
Header::_PrintGUID(const guid_t &id)
{
	static char guid[48];
	snprintf(guid, sizeof(guid),
		"%08" B_PRIx32 "-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
		B_LENDIAN_TO_HOST_INT32(id.data1), B_LENDIAN_TO_HOST_INT16(id.data2),
		B_LENDIAN_TO_HOST_INT16(id.data3), id.data4[0], id.data4[1],
		id.data4[2], id.data4[3], id.data4[4], id.data4[5], id.data4[6],
		id.data4[7]);
	return guid;
}


void
Header::_Dump(const gpt_table_header& header)
{
	dprintf("EFI header: %.8s\n", header.header);
	dprintf("EFI revision: %" B_PRIx32 "\n", header.Revision());
	dprintf("header size: %" B_PRId32 "\n", header.HeaderSize());
	dprintf("header CRC: %" B_PRIx32 "\n", header.HeaderCRC());
	dprintf("absolute block: %" B_PRIu64 "\n", header.AbsoluteBlock());
	dprintf("alternate block: %" B_PRIu64 "\n", header.AlternateBlock());
	dprintf("first usable block: %" B_PRIu64 "\n", header.FirstUsableBlock());
	dprintf("last usable block: %" B_PRIu64 "\n", header.LastUsableBlock());
	dprintf("disk GUID: %s\n", _PrintGUID(header.disk_guid));
	dprintf("entries block: %" B_PRIu64 "\n", header.EntriesBlock());
	dprintf("entry size:  %" B_PRIu32 "\n", header.EntrySize());
	dprintf("entry count: %" B_PRIu32 "\n", header.EntryCount());
	dprintf("entries CRC: %" B_PRIx32 "\n", header.EntriesCRC());
}


void
Header::_DumpPartitions()
{
	for (uint32 i = 0; i < EntryCount(); i++) {
		const gpt_partition_entry &entry = EntryAt(i);

		if (entry.partition_type == kEmptyGUID)
			continue;

		dprintf("[%3" B_PRIu32 "] partition type: %s\n", i,
			_PrintGUID(entry.partition_type));
		dprintf("      unique id: %s\n", _PrintGUID(entry.unique_guid));
		dprintf("      start block: %" B_PRIu64 "\n", entry.StartBlock());
		dprintf("      end block: %" B_PRIu64 "\n", entry.EndBlock());
		dprintf("      size: %g MB\n", (entry.EndBlock() - entry.StartBlock())
			* 512 / 1024.0 / 1024.0);
		dprintf("      attributes: %" B_PRIx64 "\n", entry.Attributes());

		char name[64];
		to_utf8(entry.name, EFI_PARTITION_NAME_LENGTH, name, sizeof(name));
		dprintf("      name: %s\n", name);
	}
}
#endif	// TRACE_EFI_GPT


}	// namespace EFI