[PATCH 7/7] LoongArch: Add pe_hdr->machine check for pei format images

[Youling Tang]

From: Youling Tang <tangyouling at kylinios.cn>

Add the pe_hdr->machine check to ensure that the kernel image is of the 64-bit
LoongArch architecture.

Signed-off-by: Youling Tang <tangyouling at kylinios.cn>
---
 kexec/arch/loongarch/image-header.h        |  3 +++
 kexec/arch/loongarch/kexec-pei-loongarch.c | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/kexec/arch/loongarch/image-header.h b/kexec/arch/loongarch/image-header.h
index 0db7615..fc71ff2 100644
--- a/kexec/arch/loongarch/image-header.h
+++ b/kexec/arch/loongarch/image-header.h
@@ -35,6 +35,9 @@ struct loongarch_image_header {
 static const uint8_t loongarch_image_pe_sig[2] = {'M', 'Z'};
 static const uint8_t loongarch_pe_machtype[6] = {'P','E', 0x0, 0x0, 0x64, 0x62};
 
+#define IMAGE_FILE_MACHINE_LOONGARCH64  0x6264 /* LoongArch 64-bit processor family */
+
+
 /**
  * loongarch_header_check_pe_sig - Helper to check the loongarch image header.
  *
diff --git a/kexec/arch/loongarch/kexec-pei-loongarch.c b/kexec/arch/loongarch/kexec-pei-loongarch.c
index 1a19a39..e0a82b6 100644
--- a/kexec/arch/loongarch/kexec-pei-loongarch.c
+++ b/kexec/arch/loongarch/kexec-pei-loongarch.c
@@ -24,9 +24,20 @@
 #include "kexec-loongarch.h"
 #include "arch/options.h"
 
+#include <pe.h>
+
+static inline int loongarch_pe_check_machine(const struct pe_hdr *pe_hdr)
+{
+	if (!pe_hdr)
+		return 0;
+
+	return (pe_hdr->machine == IMAGE_FILE_MACHINE_LOONGARCH64);
+}
+
 int pei_loongarch_probe(const char *kernel_buf, off_t kernel_size)
 {
 	const struct loongarch_image_header *h;
+	const struct pe_hdr *pe_hdr;
 
 	if (kernel_size < sizeof(struct loongarch_image_header)) {
 		dbgprintf("%s: No loongarch image header.\n", __func__);
@@ -40,6 +51,12 @@ int pei_loongarch_probe(const char *kernel_buf, off_t kernel_size)
 		return -1;
 	}
 
+	pe_hdr = (const struct pe_hdr *)(kernel_buf + get_pehdr_offset(kernel_buf));
+	if (!loongarch_pe_check_machine(pe_hdr)) {
+		dbgprintf("%s: Bad loongarch pe_hdr machine.\n", __func__);
+		return -1;
+	}
+
 	return 0;
 }
 
-- 
2.34.1