3 files changed, 74 insertions, 5 deletions
diff --git a/EfiGuardDxe/EfiGuardDxe.c b/EfiGuardDxe/EfiGuardDxe.c
index dc0b5fe..81cfb00 100644
--- a/EfiGuardDxe/EfiGuardDxe.c
+++ b/EfiGuardDxe/EfiGuardDxe.c
@@ -104,6 +104,10 @@ SetServicePointer(
 	ASSERT(gBS->CalculateCrc32 != NULL);
 
 	CONST EFI_TPL Tpl = gBS->RaiseTPL(TPL_HIGH_LEVEL); // Note: implies cli
+	CONST UINTN Cr0 = AsmReadCr0();
+	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
+	if (WpSet)
+		AsmWriteCr0(Cr0 & ~CR0_WP);
 
 	VOID* OriginalFunction = InterlockedCompareExchangePointer(ServiceTableFunction,
 																*ServiceTableFunction,
@@ -113,6 +117,8 @@ SetServicePointer(
 	ServiceTableHeader->CRC32 = 0;
 	gBS->CalculateCrc32((UINT8*)ServiceTableHeader, ServiceTableHeader->HeaderSize, &ServiceTableHeader->CRC32);
 
+	if (WpSet)
+		AsmWriteCr0(Cr0);
 	gBS->RestoreTPL(Tpl);
 
 	return OriginalFunction;
@@ -632,11 +638,6 @@ EfiGuardInitialize(
 	gKernelPatchInfo.BuildNumber = 0;
 	gKernelPatchInfo.KernelBase = NULL;
 
-	// Wipe our image info and PE headers
-	LocalImageInfo->DeviceHandle = LocalImageInfo->FilePath = LocalImageInfo->ParentHandle = NULL;
-	CONST PEFI_IMAGE_NT_HEADERS NtHeaders = RtlpImageNtHeaderEx(LocalImageInfo->ImageBase, LocalImageInfo->ImageSize);
-	ZeroMem(LocalImageInfo->ImageBase, NtHeaders->OptionalHeader.SizeOfHeaders);
-
 	// The ASCII banner is very pretty - ensure the user has enough time to admire it
 	RtlSleep(1500);
 
diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index c1ad213..4c87f89 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -2,6 +2,7 @@
 #include "util.h"
 
 #include <Library/UefiLib.h>
+#include <Library/BaseMemoryLib.h>
 #include <Library/DevicePathLib.h>
 #include <Library/PrintLib.h>
 #include <Library/UefiBootServicesTableLib.h>
@@ -99,6 +100,48 @@ PrintKernelPatchInfo(
 	}
 }
 
+VOID*
+EFIAPI
+CopyWpMem(
+	OUT VOID *Destination,
+	IN CONST VOID *Source,
+	IN UINTN Length
+	)
+{
+	CONST UINTN Cr0 = AsmReadCr0();
+	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
+	if (WpSet)
+		AsmWriteCr0(Cr0 & ~CR0_WP);
+
+	VOID* Result = CopyMem(Destination, Source, Length);
+	
+	if (WpSet)
+		AsmWriteCr0(Cr0);
+
+	return Result;
+}
+
+VOID*
+EFIAPI
+SetWpMem(
+	OUT VOID *Destination,
+	IN UINTN Length,
+	IN UINT8 Value
+	)
+{
+	CONST UINTN Cr0 = AsmReadCr0();
+	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
+	if (WpSet)
+		AsmWriteCr0(Cr0 & ~CR0_WP);
+
+	VOID* Result = SetMem(Destination, Length, Value);
+	
+	if (WpSet)
+		AsmWriteCr0(Cr0);
+
+	return Result;
+}
+
 INTN
 EFIAPI
 StrniCmp(
diff --git a/EfiGuardDxe/util.h b/EfiGuardDxe/util.h
index d69fc89..89c13e6 100644
--- a/EfiGuardDxe/util.h
+++ b/EfiGuardDxe/util.h
@@ -6,6 +6,9 @@
 #include <Zydis/Formatter.h>
 #endif
 
+#define CR0_WP		((UINTN)0x10000) // CR0.WP
+
+
 //
 // Stalls CPU for N milliseconds
 //
@@ -48,6 +51,28 @@ PrintKernelPatchInfo(
 	);
 
 //
+// Wrapper for CopyMem() that disables write protection prior to copying if needed.
+//
+VOID*
+EFIAPI
+CopyWpMem(
+	OUT VOID *Destination,
+	IN CONST VOID *Source,
+	IN UINTN Length
+	);
+
+//
+// Wrapper for SetMem() that disables write protection prior to copying if needed.
+//
+VOID*
+EFIAPI
+SetWpMem(
+	OUT VOID *Destination,
+	IN UINTN Length,
+	IN UINT8 Value
+	);
+
+//
 // Case-insensitive string comparison.
 //
 INTN