mbr_gpt.cpp

/*
Author: Blake Warner
License: public domain
*/

class Mbr
{
public:
	class Partition
	{
	public:
		enum class Type : u8
		{
			EMPTY = 0x00,
			FAT12 = 0x01,
			FAT16_SMALL = 0x04,
			FAT16_LARGE = 0x06,
			FAT32 = 0x1B,
			EXFAT = 0x07,
			EXTENDED = 0x05,
			EXTENDED_LONG = 0x0F,
			GPT = 0xEE,
			UEFI = 0xEF
		};
		//protected:
		u8 m_status;
		u24 m_chsFirstAbsoluteSector;
		Type m_partitionType;
		u24 m_chsLastAbsoluteSector;
		u32 m_lbaFirstAbsoluteSector;
		u32 m_sectorCount;
	} PACKED;

	static const u8 PARTITION_ENTRY_MAX = 4;
	//protected:
	u8 m_bootstrap[218];

	u16 m_padding;
	u8 m_originalPhysicalDrive;

	u8 m_seconds;
	u8 m_minutes;
	u8 m_hours;

	u8 m_bootstrap2[216];
	u32 m_diskSignature;
	u16 m_copyProtection;

	Partition m_partitionEntires[PARTITION_ENTRY_MAX];

	u16 m_bootSignature;

	bool isCopyProtected() const
	{
		return m_copyProtection == 0x5A5A;
	}
} PACKED;

static_assert(sizeof(Mbr) == 512, "mbr incorrect size");
static_assert(sizeof(Mbr::Partition) == 16, "mbr partition incorrect size");


	class Gpt
	{
	public:
		static const u64 MAGIC = 0x5452415020494645;

		class Partition
		{
		public:
			string name() const;
			const u64& firstLba() const { return m_firstLba; }
			const u64& lastLba() const { return m_lastLba; }
		protected:
			u8 m_typeGuid[16];
			u8 m_guid[16];
			u64 m_firstLba;
			u64 m_lastLba;
			u64 m_attributes;
			u16 m_typeName[36];
		};

		static void parse(sptr<Base>& storage, const Mbr::Partition& entry);
	protected:
		u64 m_magic;
		u32 m_revision;
		u32 m_headerSize;
		u32 m_headerCrc32;
		u32 m_padding;
		u64 m_currentLba;
		u64 m_backupLba;
		u64 m_firstUsableLba;
		u64 m_lastUsableLba;
		u8 m_diskGuid[16];
		u64 m_partitionEntriesStartLba;
		u32 m_partitionEntryCount;
		u32 m_partitionEntrySize;
		u32 m_parititionArray;
	} PACKED;

	static_assert(sizeof(Gpt) == 92, "gpt incorrect size");


	string Gpt::Partition::name() const
	{
		string r;

		for(int i = 0; i < sizeof(m_typeName) / sizeof(u16); i++)
		{
			if(m_typeName[i] < 0x80)
			{
				r.push((u8)m_typeName[i]);
			}
		}
		r.push(0);
		return r;
	}

	void Gpt::parse(sptr<Base>& storage, const Mbr::Partition& entry)
	{
		if(!storage)
		{
			log::Warning() << EncryptedString("GPT storage empty!");
			return;
		}

		Vector<u8> buffer;
		storage->readSector(buffer, 1, entry.m_lbaFirstAbsoluteSector);

		if(buffer.size() < sizeof(Gpt))
		{
			log::Warning() << EncryptedString("Error reading GPT sector");
			return;
		}

		Gpt gpt = *(Gpt*)buffer.buffer();

		if(gpt.m_magic != Gpt::MAGIC)
		{
			log::Warning() << EncryptedString("Invalid GPT magic");
			return;
		}

		storage->readSector(buffer, gpt.m_partitionEntrySize * gpt.m_partitionEntryCount / storage->sectorSize() + 1, gpt.m_partitionEntriesStartLba);

		const Gpt::Partition* partitions = (Gpt::Partition*)buffer.buffer();

		log::Debug() << EncryptedString("Found ") << gpt.m_partitionEntryCount << EncryptedString(" GPT partitions");

		for(u32 i = 0; i < gpt.m_partitionEntryCount; i++)
		{
			const Gpt::Partition& partition = partitions[i];
			if(partition.firstLba() || partition.lastLba())
			{
				auto p = std::make_shared<storage::Partition>(partition.name(), partition.firstLba(), partition.lastLba() - partition.firstLba() + 1, storage);
				storage->addPartition(p);
			}
		}

		return;
	}


	bool Storage::loadPartitions(sptr<Base>& storage)
	{
		if(!storage)
		{
			log::Debug() <<  tr("empty storage");
			return false;
		}
		Vector<u8> buffer;
		storage->readSector(buffer, 1, 0);

		if(buffer.size() < sizeof(Mbr))
		{
			log::Warning() <<  tr("mbr read failed");
			return false;
		}

		storage->m_mbr = *(Mbr*)buffer.buffer();

		for(int i = 0; i < Mbr::PARTITION_ENTRY_MAX; i++)
		{
			const auto& e = storage->m_mbr.m_partitionEntires[i];

			switch(e.m_partitionType)
			{
				case Mbr::Partition::Type::EMPTY:
					log::Debug() <<  tr("Skipping empty MBR partition");
					break;
				case Mbr::Partition::Type::GPT:
					Gpt::parse(storage, e);
					break;
				case Mbr::Partition::Type::EXTENDED:
					break;
				case Mbr::Partition::Type::EXTENDED_LONG:
					break;
				case Mbr::Partition::Type::FAT12:
				case Mbr::Partition::Type::FAT16_SMALL:
				case Mbr::Partition::Type::FAT16_LARGE:
				case Mbr::Partition::Type::FAT32:
				case Mbr::Partition::Type::EXFAT:
				{
					log::Debug() <<  tr("Found MBR partition");
					auto p = std::make_shared<Partition>("UNKNOWN", e.m_lbaFirstAbsoluteSector, e.m_sectorCount, storage);
					storage->addPartition(p);
					break;
				}
				default:
					log::Debug() <<  tr("unknown MBR partition: ") << HX(e.m_partitionType).c_str();
					auto p = std::make_shared<Partition>("UNKNOWN", e.m_lbaFirstAbsoluteSector, e.m_sectorCount, storage);
					storage->addPartition(p);
			}
		}

		if(storage->partitions().size())
		{
			log::Debug() << "loading partitions";
			for(auto& p : storage->partitions())
			{
				scanHdd(p->scheme().str() + ":/");
			}
		}

		return true;
	}

	bool Storage::addPartition(sptr<Partition>& partition)
	{
		log::Debug() <<  tr("found partition ") << partition->label().c_str();
		if(registerPartition(partition))
		{
			m_partitions.push(partition);
			return true;
		}
		return false;
	}