1 files changed, 59 insertions, 1 deletions

diff --git a/post-process-pe.c b/post-process-pe.c
index de8f4a38..008da93b 100644
--- a/post-process-pe.c
+++ b/post-process-pe.c
@@ -43,6 +43,7 @@ static int verbosity;
 	})
 
 static bool set_nx_compat = false;
+static bool require_nx_compat = false;
 
 typedef uint8_t UINT8;
 typedef uint16_t UINT16;
@@ -162,6 +163,7 @@ load_pe(const char *const file, void *const data, const size_t datasize,
 		ctx->ImageSize = PEHdr->Pe32Plus.OptionalHeader.SizeOfImage;
 		ctx->SectionAlignment =
 			PEHdr->Pe32Plus.OptionalHeader.SectionAlignment;
+		ctx->DllCharacteristics = PEHdr->Pe32Plus.OptionalHeader.DllCharacteristics;
 		FileAlignment = PEHdr->Pe32Plus.OptionalHeader.FileAlignment;
 		OptHeaderSize = sizeof(EFI_IMAGE_OPTIONAL_HEADER64);
 	} else {
@@ -172,6 +174,7 @@ load_pe(const char *const file, void *const data, const size_t datasize,
 		ctx->ImageSize = (UINT64)PEHdr->Pe32.OptionalHeader.SizeOfImage;
 		ctx->SectionAlignment =
 			PEHdr->Pe32.OptionalHeader.SectionAlignment;
+		ctx->DllCharacteristics = PEHdr->Pe32.OptionalHeader.DllCharacteristics;
 		FileAlignment = PEHdr->Pe32.OptionalHeader.FileAlignment;
 		OptHeaderSize = sizeof(EFI_IMAGE_OPTIONAL_HEADER32);
 	}
@@ -358,6 +361,50 @@ set_dll_characteristics(PE_COFF_LOADER_IMAGE_CONTEXT *ctx)
 	} else {
 		ctx->PEHdr->Pe32.OptionalHeader.DllCharacteristics = newflags;
 	}
+	ctx->DllCharacteristics = newflags;
+}
+
+static int
+validate_nx_compat(PE_COFF_LOADER_IMAGE_CONTEXT *ctx)
+{
+	EFI_IMAGE_SECTION_HEADER *Section;
+	int i;
+	bool nx_compat_flag;
+	const int level = require_nx_compat ? ERROR : WARNING;
+	int ret = 0;
+
+	nx_compat_flag = EFI_IMAGE_DLLCHARACTERISTICS_NX_COMPAT
+			 & ctx->DllCharacteristics;
+	debug(NOISE, "NX-Compat-Flag: %s\n", nx_compat_flag ? "set" : "unset");
+	if (!nx_compat_flag) {
+		debug(level, "NX Compatibility flag is not set\n");
+		if (require_nx_compat)
+			ret = -1;
+	}
+
+	debug(NOISE, "Section alignment is 0x%x, page size is 0x%x\n",
+	      ctx->SectionAlignment, PAGE_SIZE);
+	if (ctx->SectionAlignment != PAGE_SIZE) {
+		debug(level, "Section alignment is not page aligned\n");
+		if (require_nx_compat)
+			ret = -1;
+	}
+
+	Section = ctx->FirstSection;
+	for (i=0, Section = ctx->FirstSection; i < ctx->NumberOfSections; i++, Section++) {
+		debug(NOISE, "Section %d has WRITE=%d and EXECUTE=%d\n", i,
+		      (Section->Characteristics & EFI_IMAGE_SCN_MEM_WRITE) ? 1 : 0,
+		      (Section->Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) ? 1 : 0);
+
+		if ((Section->Characteristics & EFI_IMAGE_SCN_MEM_WRITE) &&
+		    (Section->Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE)) {
+			debug(level, "Section %d is writable and executable\n", i);
+			if (require_nx_compat)
+				ret = -1;
+		}
+	}
+
+	return ret;
 }
 
 static void
@@ -449,6 +496,10 @@ handle_one(char *f)
 
 	set_dll_characteristics(&ctx);
 
+	rc = validate_nx_compat(&ctx);
+	if (rc < 0)
+		err(2, "NX compatibility check failed\n");
+
 	fix_timestamp(&ctx);
 
 	fix_checksum(&ctx, map, sz);
@@ -483,6 +534,7 @@ static void __attribute__((__noreturn__)) usage(int status)
 	fprintf(out, "       -v    Be more verbose\n");
 	fprintf(out, "       -N    Disable the NX compatibility flag\n");
 	fprintf(out, "       -n    Enable the NX compatibility flag\n");
+	fprintf(out, "       -x    Error on NX incompatibility\n");
 	fprintf(out, "       -h    Print this help text and exit\n");
 
 	exit(status);
@@ -510,11 +562,14 @@ int main(int argc, char **argv)
 		{.name = "verbose",
 		 .val = 'v',
 		},
+		{.name = "error-nx-compat",
+		 .val = 'x',
+		},
 		{.name = ""}
 	};
 	int longindex = -1;
 
-	while ((i = getopt_long(argc, argv, "hNnqv", options, &longindex)) != -1) {
+	while ((i = getopt_long(argc, argv, "hNnqvx", options, &longindex)) != -1) {
 		switch (i) {
 		case 'h':
 		case '?':
@@ -532,6 +587,9 @@ int main(int argc, char **argv)
 		case 'v':
 			verbosity = MIN(verbosity + 1, MAX_VERBOSITY);
 			break;
+		case 'x':
+			require_nx_compat = true;
+			break;
 		}
 	}