Bootloader: installer now patches GPT GUID's

bootloader/install.sh

@@ -7,6 +7,10 @@ if [[ -z $BANAN_DISK_IMAGE_PATH ]]; then
 	exit 1
 fi
 
+ROOT_PARTITION_INDEX=2
+ROOT_PARTITION_INFO=$(fdisk -x $BANAN_DISK_IMAGE_PATH | grep "^$BANAN_DISK_IMAGE_PATH" | head -$ROOT_PARTITION_INDEX | tail -1)
+ROOT_PARTITION_GUID=$(echo $ROOT_PARTITION_INFO | cut -d' ' -f6)
+
 CURRENT_DIR=$(dirname $(realpath $0))
 
 INSTALLER_DIR=$CURRENT_DIR/installer
@@ -31,5 +35,5 @@ x86_64-banan_os-as $CURRENT_DIR/arch/x86_64/boot.S -o $BUILD_DIR/bootloader.o
 echo linking bootloader
 x86_64-banan_os-ld -nostdlib -T $CURRENT_DIR/arch/x86_64/linker.ld $BUILD_DIR/bootloader.o -o $BUILD_DIR/bootloader
 
-echo installing bootloader
+echo installing bootloader to
-$INSTALLER_BUILD_DIR/x86_64-banan_os-bootloader-installer $BUILD_DIR/bootloader $BANAN_DISK_IMAGE_PATH
+$INSTALLER_BUILD_DIR/x86_64-banan_os-bootloader-installer $BUILD_DIR/bootloader $BANAN_DISK_IMAGE_PATH $ROOT_PARTITION_GUID

bootloader/installer/GPT.cpp

@@ -63,18 +63,18 @@ const GPTHeader& GPTFile::gpt_header() const
     return *reinterpret_cast<GPTHeader*>(m_mmap + SECTOR_SIZE);
 }
 
-bool GPTFile::install_bootcode(std::span<const uint8_t> boot_code)
+bool GPTFile::install_stage1(std::span<const uint8_t> stage1)
 {
 	auto& mbr = this->mbr();
 
-	if (boot_code.size() > sizeof(mbr.boot_code))
+	if (stage1.size() > sizeof(mbr.boot_code))
 	{
-		std::cerr << m_path << ": can't fit " << boot_code.size() << " bytes of boot code in mbr (max is " << sizeof(mbr.boot_code) << ")" << std::endl;
+		std::cerr << m_path << ": can't fit " << stage1.size() << " bytes of boot code in mbr (max is " << sizeof(mbr.boot_code) << ")" << std::endl;
 		return false;
 	}
 
 	// copy boot code
-	memcpy(mbr.boot_code, boot_code.data(), boot_code.size());
+	memcpy(mbr.boot_code, stage1.data(), stage1.size());
 
 	// setup mbr
 	mbr.unique_mbr_disk_signature = 0xdeadbeef;
@@ -99,29 +99,104 @@ bool GPTFile::install_bootcode(std::span<const uint8_t> boot_code)
 	return true;
 }
 
-bool GPTFile::write_partition(std::span<const uint8_t> data, GUID type_guid)
+bool GPTFile::install_stage2(std::span<const uint8_t> stage2, const GUID& root_partition_guid)
 {
-	auto partition = find_partition(type_guid);
+	if (stage2.size() < 16)
+	{
+		std::cerr << m_path << ": contains invalid .stage2 section, too small for patches" << std::endl;
+		return false;
+	}
+
+	// find GUID patch offsets
+	std::size_t disk_guid_offset(-1);
+	std::size_t part_guid_offset(-1);
+	for (std::size_t i = 0; i < stage2.size() - 16; i++)
+	{
+		if (memcmp(stage2.data() + i, "root disk guid  ", 16) == 0)
+		{
+			if (disk_guid_offset != std::size_t(-1))
+			{
+				std::cerr << m_path << ": contains invalid .stage2 section, multiple patchable disk guids" << std::endl;
+				return false;
+			}
+			disk_guid_offset = i;
+		}
+		if (memcmp(stage2.data() + i, "root part guid  ", 16) == 0)
+		{
+			if (part_guid_offset != std::size_t(-1))
+			{
+				std::cerr << m_path << ": contains invalid .stage2 section, multiple patchable partition guids" << std::endl;
+				return false;
+			}
+			part_guid_offset = i;
+		}
+	}
+	if (disk_guid_offset == std::size_t(-1))
+	{
+		std::cerr << m_path << ": contains invalid .stage2 section, no patchable disk guid" << std::endl;
+		return false;
+	}
+	if (part_guid_offset == std::size_t(-1))
+	{
+		std::cerr << m_path << ": contains invalid .stage2 section, no patchable partition guid" << std::endl;
+		return false;
+	}
+	
+
+	auto partition = find_partition_with_type(bios_boot_guid);
 	if (!partition.has_value())
 	{
-		std::cerr << m_path << ": could not find partition with type " << type_guid << std::endl;
+		std::cerr << m_path << ": could not find partition with type " << bios_boot_guid << std::endl;
 		return false;
 	}
 
 	const std::size_t partition_size = (partition->ending_lba - partition->starting_lba + 1) * SECTOR_SIZE;
 
-	if (data.size() > partition_size)
+	if (stage2.size() > partition_size)
 	{
-		std::cerr << m_path << ": can't fit " << data.size() << " bytes of data to partition of size " << partition_size << std::endl;
+		std::cerr << m_path << ": can't fit " << stage2.size() << " bytes of data to partition of size " << partition_size << std::endl;
 		return false;
 	}
 
-	memcpy(m_mmap + partition->starting_lba * SECTOR_SIZE, data.data(), data.size());
+	uint8_t* partition_start = m_mmap + partition->starting_lba * SECTOR_SIZE;
+	memcpy(partition_start, stage2.data(), stage2.size());
+
+	// patch GUIDs
+	*reinterpret_cast<GUID*>(partition_start + disk_guid_offset) = gpt_header().disk_guid;
+	*reinterpret_cast<GUID*>(partition_start + part_guid_offset) = root_partition_guid;
 
 	return true;
 }
 
-std::optional<GPTPartitionEntry> GPTFile::find_partition(const GUID& type_guid) const
+bool GPTFile::install_bootloader(std::span<const uint8_t> stage1, std::span<const uint8_t> stage2, const GUID& root_partition_guid)
+{
+	if (!find_partition_with_guid(root_partition_guid).has_value())
+	{
+		std::cerr << m_path << ": no partition with GUID " << root_partition_guid << std::endl;
+		return false;
+	}
+	if (!install_stage1(stage1))
+		return false;
+	if (!install_stage2(stage2, root_partition_guid))
+		return false;
+	return true;
+}
+
+std::optional<GPTPartitionEntry> GPTFile::find_partition_with_guid(const GUID& guid) const
+{
+	const auto& gpt_header = this->gpt_header();
+	const uint8_t* partition_entry_array_start = m_mmap + gpt_header.partition_entry_lba * SECTOR_SIZE;
+	for (std::size_t i = 0; i < gpt_header.number_of_partition_entries; i++)
+	{
+		const auto& partition_entry = *reinterpret_cast<const GPTPartitionEntry*>(partition_entry_array_start + i * gpt_header.size_of_partition_entry);
+		if (partition_entry.partition_guid != guid)
+			continue;
+		return partition_entry;
+	}
+	return {};
+}
+
+std::optional<GPTPartitionEntry> GPTFile::find_partition_with_type(const GUID& type_guid) const
 {
 	const auto& gpt_header = this->gpt_header();
 	const uint8_t* partition_entry_array_start = m_mmap + gpt_header.partition_entry_lba * SECTOR_SIZE;

bootloader/installer/GPT.h

@@ -65,8 +65,7 @@ public:
 	GPTFile(std::string_view path);
 	~GPTFile();
 
-	bool install_bootcode(std::span<const uint8_t>);
+	bool install_bootloader(std::span<const uint8_t> stage1, std::span<const uint8_t> stage2, const GUID& root_partition_guid);
-	bool write_partition(std::span<const uint8_t>, GUID type_guid);
 
 	const GPTHeader& gpt_header() const;
 
@@ -77,7 +76,11 @@ public:
 private:
 	MBR& mbr();
 	bool validate_gpt_header() const;
-	std::optional<GPTPartitionEntry> find_partition(const GUID& type_guid) const;
+	std::optional<GPTPartitionEntry> find_partition_with_guid(const GUID& guid) const;
+	std::optional<GPTPartitionEntry> find_partition_with_type(const GUID& type_guid) const;
+
+	bool install_stage1(std::span<const uint8_t> stage1);
+	bool install_stage2(std::span<const uint8_t> stage2, const GUID& root_partition_guid);
 
 private:
 	const std::string	m_path;

bootloader/installer/main.cpp

@@ -7,9 +7,16 @@ int main(int argc, char** argv)
 {
 	using namespace std::string_view_literals;
 
-	if (argc != 3)
+	if (argc != 4)
 	{
-		std::fprintf(stderr, "usage: %s BOOTLOADER DISK_IMAGE}\n", argv[0]);
+		std::fprintf(stderr, "usage: %s BOOTLOADER DISK_IMAGE ROOT_PARTITION_GUID\n", argv[0]);
+		return 1;
+	}
+
+	auto root_partition_guid = GUID::from_string(argv[3]);
+	if (!root_partition_guid.has_value())
+	{
+		std::cerr << "invalid guid '" << argv[3] << '\'' << std::endl;
 		return 1;
 	}
 
@@ -29,13 +36,9 @@ int main(int argc, char** argv)
 	if (!disk_image.success())
 		return 1;
 
-	if (!disk_image.install_bootcode(*stage1))
+	if (!disk_image.install_bootloader(*stage1, *stage2, *root_partition_guid))
 		return 1;
-	std::cout << "wrote stage1 bootloader" << std::endl;
+	std::cout << "bootloader installed" << std::endl;
-
-	if (!disk_image.write_partition(*stage2, bios_boot_guid))
-		return 1;
-	std::cout << "wrote stage2 bootloader" << std::endl;
 
 	return 0;
 }