Fix ACPI table handling and update the main logic

Apparently the BGRT must be stored in every XSDT on the system.
Also some table checksums were not updated correctly.
This patch cleans up the whole process of updating the ACPI tables.

Makefile

@@ -7,7 +7,7 @@ LIBS    = -L$(GNUEFI_LIB) -lefi -lgcc
 GNUEFI_INC = /usr/$(CC_PREFIX)/include/efi
 GNUEFI_LIB = /usr/$(CC_PREFIX)/lib
 
-FILES_C = src/main.c src/util.c src/config.c
+FILES_C = src/main.c src/util.c src/types.c src/config.c
 FILES_H = $(wildcard src/*.h)
 
 .PHONY: all

src/config.c

@@ -65,7 +65,7 @@ static void ReadConfigImage(struct HackBGRT_config* config, const CHAR16* line)
 	} else if (StrStr(line, L"remove")) {
 		action = HackBGRT_REMOVE;
 	} else if (StrStr(line, L"black")) {
-		action = HackBGRT_BLACK;
+		action = HackBGRT_REPLACE;
 	} else if (StrStr(line, L"keep")) {
 		action = HackBGRT_KEEP;
 	} else {

src/config.h

@@ -6,7 +6,7 @@
  * Possible actions to perform on the BGRT.
  */
 enum HackBGRT_action {
-	HackBGRT_KEEP = 0, HackBGRT_REPLACE, HackBGRT_REMOVE, HackBGRT_BLACK
+	HackBGRT_KEEP = 0, HackBGRT_REPLACE, HackBGRT_REMOVE
 };
 
 /**

src/main.c

@@ -53,106 +53,77 @@ static int SelectCoordinate(int value, int automatic, int native) {
 }
 
 /**
- * Initialize (clear) a BGRT.
+ * Create a new XSDT with the given number of entries.
  *
- * @param bgrt The BGRT to initialize.
+ * @param xsdt0 The old XSDT.
+ * @param entries The number of SDT entries.
+ * @return Pointer to a new XSDT.
  */
-static void InitBGRT(ACPI_BGRT* bgrt) {
-	const char data[0x38] = "BGRT" "\x38\x00\x00\x00" "\x00" "\xd6" "Mtblx*" "HackBGRT" "\x20\x17\x00\x00" "PTL " "\x02\x00\x00\x00" "\x01\x00" "\x00" "\x00";
-	CopyMem(bgrt, data, sizeof(data));
+ACPI_SDT_HEADER* CreateXsdt(ACPI_SDT_HEADER* xsdt0, UINTN entries) {
+	ACPI_SDT_HEADER* xsdt = 0;
+	UINT32 xsdt_len = sizeof(ACPI_SDT_HEADER) + entries * sizeof(UINT64);
+	BS->AllocatePool(EfiACPIReclaimMemory, xsdt_len, (void**)&xsdt);
+	if (!xsdt) {
+		Print(L"HackBGRT: Failed to allocate memory for XSDT.\n");
+		return 0;
+	}
+	ZeroMem(xsdt, xsdt_len);
+	CopyMem(xsdt, xsdt0, min(xsdt0->length, xsdt_len));
+	xsdt->length = xsdt_len;
+	SetAcpiSdtChecksum(xsdt);
+	return xsdt;
 }
 
 /**
- * Fill a BGRT as specified by the parameters.
+ * Update the ACPI tables as needed for the desired BGRT change.
  *
- * @param bgrt The BGRT to fill.
- * @param new_bmp The BMP to use.
- * @param new_x The x coordinate to use.
- * @param new_y The y coordinate to use.
- */
-static void FillBGRT(ACPI_BGRT* bgrt, BMP* new_bmp, int new_x, int new_y) {
-	BMP* old_bmp = (BMP*) (UINTN) bgrt->image_address;
-	ACPI_BGRT bgrt0 = *bgrt;
-	InitBGRT(bgrt);
-
-	if (new_bmp) {
-		bgrt->image_address = (UINTN) new_bmp;
-	}
-	BMP* bmp = (BMP*) (UINTN) bgrt->image_address;
-
-	// Calculate the automatically centered position for the image.
-	int x_auto, y_auto;
-	if (GOP()) {
-		x_auto = max(0, ((int)GOP()->Mode->Info->HorizontalResolution - (int)bmp->width) / 2);
-		y_auto = max(0, ((int)GOP()->Mode->Info->VerticalResolution * 2/3 - (int)bmp->height) / 2);
-	} else {
-		x_auto = max(0, (int)bgrt0.image_offset_x + ((int)old_bmp->width - (int)bmp->width) / 2);
-		y_auto = max(0, (int)bgrt0.image_offset_y + ((int)old_bmp->height - (int)bmp->height) / 2);
-	}
-
-	// Set the position (manual, automatic, original).
-	bgrt->image_offset_x = SelectCoordinate(new_x, x_auto, bgrt0.image_offset_x);
-	bgrt->image_offset_y = SelectCoordinate(new_y, y_auto, bgrt0.image_offset_y);
-	Debug(L"HackBGRT: BMP at (%d, %d).\n", (int) bgrt->image_offset_x, (int) bgrt->image_offset_y);
-
-	bgrt->header.checksum = 0;
-	bgrt->header.checksum = CalculateAcpiChecksum(bgrt, sizeof(*bgrt));
-}
-
-/**
- * Find the BGRT and optionally destroy it or create if missing.
+ * If action is REMOVE, all BGRT entries will be removed.
+ * If action is KEEP, the first BGRT entry will be returned.
+ * If action is REPLACE, the given BGRT entry will be stored in each XSDT.
  *
  * @param action The intended action.
+ * @param bgrt The BGRT, if action is REPLACE.
  * @return Pointer to the BGRT, or 0 if not found (or destroyed).
  */
-static ACPI_BGRT* FindBGRT(enum HackBGRT_action action) {
-	ACPI_20_RSDP* good_rsdp = 0;
-	ACPI_BGRT* bgrt = 0;
-
+static ACPI_BGRT* HandleAcpiTables(enum HackBGRT_action action, ACPI_BGRT* bgrt) {
 	for (int i = 0; i < ST->NumberOfTableEntries; i++) {
 		EFI_GUID Acpi20TableGuid = ACPI_20_TABLE_GUID;
 		EFI_GUID* vendor_guid = &ST->ConfigurationTable[i].VendorGuid;
 		if (!CompareGuid(vendor_guid, &AcpiTableGuid) && !CompareGuid(vendor_guid, &Acpi20TableGuid)) {
 			continue;
 		}
-		EFI_CONFIGURATION_TABLE *ect = &ST->ConfigurationTable[i];
-		if (CompareMem(ect->VendorTable, "RSD PTR ", 8) != 0) {
+		ACPI_20_RSDP* rsdp = (ACPI_20_RSDP *) ST->ConfigurationTable[i].VendorTable;
+		if (CompareMem(rsdp->signature, "RSD PTR ", 8) != 0 || rsdp->revision < 2 || !VerifyAcpiRsdp2Checksums(rsdp)) {
 			continue;
 		}
-		ACPI_20_RSDP* rsdp = (ACPI_20_RSDP *)ect->VendorTable;
 		Debug(L"RSDP: revision = %d, OEM ID = %s\n", rsdp->revision, TmpStr(rsdp->oem_id, 6));
 
-		if (rsdp->revision < 2) {
-			Debug(L"* XSDT: N/A (revision < 2)\n");
-			continue;
-		}
 		ACPI_SDT_HEADER* xsdt = (ACPI_SDT_HEADER *) (UINTN) rsdp->xsdt_address;
-		if (!xsdt) {
-			Debug(L"* XSDT: N/A (null)\n");
+		if (!xsdt || CompareMem(xsdt->signature, "XSDT", 4) != 0 || !VerifyAcpiSdtChecksum(xsdt)) {
+			Debug(L"* XSDT: missing or invalid\n");
 			continue;
 		}
-		if (CompareMem(xsdt->signature, "XSDT", 4) != 0) {
-			Debug(L"* XSDT: N/A (invalid signature)\n");
-			continue;
-		}
-		good_rsdp = rsdp;
 		UINT64* entry_arr = (UINT64*)&xsdt[1];
 		UINT32 entry_arr_length = (xsdt->length - sizeof(*xsdt)) / sizeof(UINT64);
 
 		Debug(L"* XSDT: OEM ID = %s, entry count = %d\n", TmpStr(xsdt->oem_id, 6), entry_arr_length);
 
+		int bgrt_count = 0;
 		for (int j = 0; j < entry_arr_length; j++) {
 			ACPI_SDT_HEADER *entry = (ACPI_SDT_HEADER *)((UINTN)entry_arr[j]);
+			if (CompareMem(entry->signature, "BGRT", 4) != 0) {
+				continue;
+			}
 			Debug(L" - ACPI table: %s, revision = %d, OEM ID = %s\n", TmpStr(entry->signature, 4), entry->revision, TmpStr(entry->oem_id, 6));
-			if (CompareMem(entry->signature, "BGRT", 4) == 0) {
-				if (!bgrt && action != HackBGRT_REMOVE) {
-					bgrt = (void*) entry;
-				} else {
-					if (bgrt) {
-						Debug(L" -> Deleting; BGRT was already found!\n");
-					} else {
-						Debug(L" -> Deleting.\n");
+			switch (action) {
+				case HackBGRT_KEEP:
+					if (!bgrt) {
+						Debug(L" -> Returning this one for later use.\n");
+						bgrt = (ACPI_BGRT*) entry;
 					}
+					break;
+				case HackBGRT_REMOVE:
+					Debug(L" -> Deleting.\n");
 					for (int k = j+1; k < entry_arr_length; ++k) {
 						entry_arr[k-1] = entry_arr[k];
 					}
@@ -160,56 +131,36 @@ static ACPI_BGRT* FindBGRT(enum HackBGRT_action action) {
 					entry_arr[entry_arr_length] = 0;
 					xsdt->length -= sizeof(entry_arr[0]);
 					--j;
-				}
+					break;
+				case HackBGRT_REPLACE:
+					Debug(L" -> Replacing.\n");
+					entry_arr[j] = (UINTN) bgrt;
 			}
+			bgrt_count += 1;
 		}
-	}
-	if (action == HackBGRT_REMOVE) {
-		return 0;
-	}
-	if (!good_rsdp) {
-		Print(L"HackBGRT: RSDP or XSDT not found.\n");
-		return 0;
-	}
-	if (!bgrt) {
-		if (action == HackBGRT_KEEP) {
-			Print(L"HackBGRT: BGRT not found.\n");
-			return 0;
+		if (!bgrt_count && action == HackBGRT_REPLACE && bgrt) {
+			Debug(L" - Adding missing BGRT.\n");
+			xsdt = CreateXsdt(xsdt, entry_arr_length + 1);
+			entry_arr = (UINT64*)&xsdt[1];
+			entry_arr[entry_arr_length++] = (UINTN) bgrt;
+			rsdp->xsdt_address = (UINTN) xsdt;
+			SetAcpiRsdp2Checksums(rsdp);
 		}
-		Debug(L"HackBGRT: BGRT not found, creating.\n");
-		ACPI_20_RSDP* rsdp = good_rsdp;
-		ACPI_SDT_HEADER* xsdt0 = (ACPI_SDT_HEADER *) (UINTN) rsdp->xsdt_address;
-		ACPI_SDT_HEADER* xsdt = 0;
-		UINT32 xsdt_len = xsdt0->length + sizeof(UINT64);
-		BS->AllocatePool(EfiACPIReclaimMemory, xsdt_len, (void**)&xsdt);
-		BS->AllocatePool(EfiACPIReclaimMemory, sizeof(*bgrt), (void**)&bgrt);
-		if (!xsdt || !bgrt) {
-			Print(L"HackBGRT: Failed to allocate memory for XSDT and BGRT.\n");
-			return 0;
-		}
-		rsdp->xsdt_address = (UINTN) xsdt;
-		CopyMem(xsdt, xsdt0, xsdt0->length);
-		*(UINT64*)((char*)xsdt + xsdt->length) = (UINTN) bgrt;
-		xsdt->length = xsdt_len;
-		InitBGRT(bgrt);
+		SetAcpiSdtChecksum(xsdt);
 	}
 	return bgrt;
 }
 
 /**
- * Load a bitmap or generate one, or return 0 if not applicable.
+ * Load a bitmap or generate a black one.
  *
- * @param action Tells what to do.
  * @param root_dir The root directory for loading a BMP.
- * @param path The BMP path within the root directory.
- * @return The loaded BMP, or 0 if not needed or not available.
+ * @param path The BMP path within the root directory; NULL for a black BMP.
+ * @return The loaded BMP, or 0 if not available.
  */
-static BMP* LoadBMP(enum HackBGRT_action action, EFI_FILE_HANDLE root_dir, const CHAR16* path) {
+static BMP* LoadBMP(EFI_FILE_HANDLE root_dir, const CHAR16* path) {
 	BMP* bmp = 0;
-	if (action == HackBGRT_KEEP || action == HackBGRT_REMOVE) {
-		return 0;
-	}
-	if (action == HackBGRT_BLACK) {
+	if (!path) {
 		BS->AllocatePool(EfiBootServicesData, 58, (void**) &bmp);
 		if (!bmp) {
 			Print(L"HackBGRT: Failed to allocate a blank BMP!\n");
@@ -226,10 +177,6 @@ static BMP* LoadBMP(enum HackBGRT_action action, EFI_FILE_HANDLE root_dir, const
 		);
 		return bmp;
 	}
-	if (!path) {
-		Print(L"HackBGRT: Missing BMP path. REPORT THIS BUG!");
-		return 0;
-	}
 	Debug(L"HackBGRT: Loading %s.\n", path);
 	bmp = LoadFile(root_dir, path, 0);
 	if (!bmp) {
@@ -240,6 +187,85 @@ static BMP* LoadBMP(enum HackBGRT_action action, EFI_FILE_HANDLE root_dir, const
 	return bmp;
 }
 
+/**
+ * The main logic for BGRT modification.
+ *
+ * @param root_dir The root directory for loading a BMP.
+ */
+void HackBgrt(EFI_FILE_HANDLE root_dir) {
+	// REMOVE: simply delete all BGRT entries.
+	if (config.action == HackBGRT_REMOVE) {
+		HandleAcpiTables(config.action, 0);
+		return;
+	}
+
+	// KEEP/REPLACE: first get the old BGRT entry.
+	ACPI_BGRT* bgrt = HandleAcpiTables(HackBGRT_KEEP, 0);
+
+	// Get the old BMP and position, if possible.
+	BMP* old_bmp = 0;
+	int old_x = 0, old_y = 0;
+	if (bgrt && VerifyAcpiSdtChecksum(bgrt)) {
+		old_bmp = (BMP*) (UINTN) bgrt->image_address;
+		old_x = bgrt->image_offset_x;
+		old_y = bgrt->image_offset_y;
+	}
+
+	// Missing BGRT?
+	if (!bgrt) {
+		// Keep missing = do nothing.
+		if (config.action == HackBGRT_KEEP) {
+			return;
+		}
+		// Replace missing = allocate new.
+		BS->AllocatePool(EfiACPIReclaimMemory, sizeof(*bgrt), (void**)&bgrt);
+		if (!bgrt) {
+			Print(L"HackBGRT: Failed to allocate memory for BGRT.\n");
+			return;
+		}
+	}
+
+	// Clear the BGRT.
+	const char data[0x38] =
+		"BGRT" "\x38\x00\x00\x00" "\x00" "\xd6" "Mtblx*" "HackBGRT"
+		"\x20\x17\x00\x00" "PTL " "\x02\x00\x00\x00"
+		"\x01\x00" "\x00" "\x00";
+	CopyMem(bgrt, data, sizeof(data));
+
+	// Get the image (either old or new).
+	BMP* new_bmp = old_bmp;
+	if (config.action == HackBGRT_REPLACE) {
+		new_bmp = LoadBMP(root_dir, config.image_path);
+	}
+
+	// No image = no need for BGRT.
+	if (!new_bmp) {
+		HandleAcpiTables(HackBGRT_REMOVE, 0);
+		return;
+	}
+
+	bgrt->image_address = (UINTN) new_bmp;
+
+	// Calculate the automatically centered position for the image.
+	int auto_x = 0, auto_y = 0;
+	if (GOP()) {
+		auto_x = max(0, ((int)GOP()->Mode->Info->HorizontalResolution - (int)new_bmp->width) / 2);
+		auto_y = max(0, ((int)GOP()->Mode->Info->VerticalResolution * 2/3 - (int)new_bmp->height) / 2);
+	} else if (old_bmp) {
+		auto_x = max(0, old_x + ((int)old_bmp->width - (int)new_bmp->width) / 2);
+		auto_y = max(0, old_y + ((int)old_bmp->height - (int)new_bmp->height) / 2);
+	}
+
+	// Set the position (manual, automatic, original).
+	bgrt->image_offset_x = SelectCoordinate(config.image_x, auto_x, old_x);
+	bgrt->image_offset_y = SelectCoordinate(config.image_y, auto_y, old_y);
+	Debug(L"HackBGRT: BMP at (%d, %d).\n", (int) bgrt->image_offset_x, (int) bgrt->image_offset_y);
+
+	// Store this BGRT in the ACPI tables.
+	SetAcpiSdtChecksum(bgrt);
+	HandleAcpiTables(HackBGRT_REPLACE, bgrt);
+}
+
 /**
  * The main program.
  */
@@ -269,24 +295,26 @@ EFI_STATUS EFIAPI EfiMain(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *ST_) {
 	}
 	Debug = config.debug ? Print : NullPrint;
 
-	BMP* new_bmp = LoadBMP(config.action, root_dir, config.image_path);
-	ACPI_BGRT* bgrt = FindBGRT(config.action);
-	if (bgrt) {
-		FillBGRT(bgrt, new_bmp, config.image_x, config.image_y);
-	}
+	HackBgrt(root_dir);
 
 	if (!config.boot_path) {
 		Print(L"HackBGRT: Boot path not specified.\n");
 		goto fail;
 	}
 
-	Debug(L"HackBGRT: Loading and booting %s.\n", config.boot_path);
+	Debug(L"HackBGRT: Loading %s.\n", config.boot_path);
 	EFI_DEVICE_PATH* boot_dp = FileDevicePath(image->DeviceHandle, (CHAR16*) config.boot_path);
 	EFI_HANDLE next_image_handle;
 	if (EFI_ERROR(BS->LoadImage(0, image_handle, boot_dp, 0, 0, &next_image_handle))) {
 		Print(L"HackBGRT: LoadImage for new image (%s) failed.\n", config.boot_path);
 		goto fail;
 	}
+	if (config.debug) {
+		Print(L"HackBGRT: Ready to boot.\nPress escape to cancel, any other key to boot.\n");
+		if (ReadKey().ScanCode == SCAN_ESC) {
+			return 0;
+		}
+	}
 	if (EFI_ERROR(BS->StartImage(next_image_handle, 0, 0))) {
 		Print(L"HackBGRT: StartImage for %s failed.\n", config.boot_path);
 		goto fail;
@@ -296,7 +324,7 @@ EFI_STATUS EFIAPI EfiMain(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *ST_) {
 
 	fail: {
 		Print(L"HackBGRT has failed. Use parameter debug=1 for details.\nPress any key to exit.\n");
-		WaitKey();
+		ReadKey();
 		return 1;
 	}
 }

src/types.c

@@ -0,0 +1,37 @@
+#include "types.h"
+
+UINT8 SumBytes(const UINT8* arr, UINTN size) {
+	UINT8 sum = 0;
+	for (UINTN i = 0; i < size; ++i) {
+		sum += arr[i];
+	}
+	return sum;
+}
+
+int VerifyAcpiRsdp2Checksums(const void* data) {
+	const UINT8* arr = data;
+	UINTN size = *(const UINT32*)&arr[20];
+	return SumBytes(arr, 20) == 0 && SumBytes(arr, size) == 0;
+}
+
+void SetAcpiRsdp2Checksums(void* data) {
+	UINT8* arr = data;
+	UINTN size = *(const UINT32*)&arr[20];
+	arr[9] = 0;
+	arr[32] = 0;
+	arr[9] = -SumBytes(arr, 20);
+	arr[32] = -SumBytes(arr, size);
+}
+
+int VerifyAcpiSdtChecksum(const void* data) {
+	const UINT8* arr = data;
+	UINTN size = *(const UINT32*)&arr[4];
+	return SumBytes(arr, size) == 0;
+}
+
+void SetAcpiSdtChecksum(void* data) {
+	UINT8* arr = data;
+	UINTN size = *(const UINT32*)&arr[4];
+	arr[9] = 0;
+	arr[9] = -SumBytes(arr, size);
+}

src/types.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <efi.h>
+
 #pragma pack(push, 1)
 
 /** RSDP (Root System Description Pointer) */
@@ -51,4 +53,35 @@ typedef struct {
 	UINT16 planes;
 	UINT16 bpp;
 } BMP;
+
+/**
+ * Verify the checksums of an ACPI RSDP version 2.
+ *
+ * @param data Pointer to the table.
+ * @return 1 if the checksum is correct, 0 otherwise.
+ */
+extern int VerifyAcpiRsdp2Checksums(const void* data);
+
+/**
+ * Set the correct checksums of an ACPI RSDP version 2.
+ *
+ * @param data Pointer to the table.
+ */
+extern void SetAcpiRsdp2Checksums(void* data);
+
+/**
+ * Verify the checksum of an ACPI SDT.
+ *
+ * @param data Pointer to the table.
+ * @return 1 if the checksum is correct, 0 otherwise.
+ */
+extern int VerifyAcpiSdtChecksum(const void* data);
+
+/**
+ * Set the correct checksum for an ACPI SDT.
+ *
+ * @param data Pointer to the table.
+ */
+extern void SetAcpiSdtChecksum(void* data);
+
 #pragma pack(pop)

src/util.c

@@ -18,15 +18,6 @@ UINTN NullPrint(IN CHAR16 *fmt, ...) {
 	return 0;
 }
 
-UINT8 CalculateAcpiChecksum(void* data, UINTN size) {
-	UINT8 sum = 0;
-	UINT8* arr = data;
-	for (UINTN i = 0; i < size; ++i) {
-		sum += arr[i];
-	}
-	return 256 - sum;
-}
-
 const CHAR16* TrimLeft(const CHAR16* s) {
 	// Skip white-space and BOM.
 	while (s[0] == L'\xfeff' || s[0] == ' ' || s[0] == '\t') {
@@ -82,6 +73,13 @@ void WaitKey(void) {
 	WaitForSingleEvent(ST->ConIn->WaitForKey, 0);
 }
 
+EFI_INPUT_KEY ReadKey(void) {
+	WaitKey();
+	EFI_INPUT_KEY key = {0};
+	ST->ConIn->ReadKeyStroke(ST->ConIn, &key);
+	return key;
+}
+
 void* LoadFileWithPadding(EFI_FILE_HANDLE dir, const CHAR16* path, UINTN* size_ptr, UINTN padding) {
 	EFI_STATUS e;
 	EFI_FILE_HANDLE handle;

src/util.h

@@ -16,15 +16,6 @@ extern const CHAR16* TmpStr(CHAR8 *src, int length);
  */
 extern UINTN NullPrint(IN CHAR16 *fmt, ...);
 
-/**
- * Calculate the checksum for an ACPI table.
- *
- * @param data Pointer to the table.
- * @param size Table length in bytes.
- * @return Checksum.
- */
-extern UINT8 CalculateAcpiChecksum(void* data, UINTN size);
-
 /**
  * Return the greater of two numbers.
  */
@@ -32,6 +23,13 @@ static inline int max(int a, int b) {
 	return a > b ? a : b;
 }
 
+/**
+ * Return the smaller of two numbers.
+ */
+static inline int min(int a, int b) {
+	return a < b ? a : b;
+}
+
 /**
  * Trim BOM, spaces and tabs from the beginning of a string.
  *
@@ -83,10 +81,17 @@ extern void RandomSeed(UINT64 a, UINT64 b);
 extern void RandomSeedAuto(void);
 
 /**
- * Wait for a key press.
+ * Wait for a key press. It will still remain in the buffer.
  */
 extern void WaitKey(void);
 
+/**
+ * Wait for a key press and read it.
+ *
+ * @return The pressed key.
+ */
+extern EFI_INPUT_KEY ReadKey(void);
+
 /**
  * Load a file, allocate some extra bytes as well.
  */