Disable CET when clearing CR0.WP

7 files changed, 153 insertions, 14 deletions

diff --git a/EfiGuardDxe/EfiGuardDxe.inf b/EfiGuardDxe/EfiGuardDxe.inf
index 67df756..2ab9719 100644
--- a/EfiGuardDxe/EfiGuardDxe.inf
+++ b/EfiGuardDxe/EfiGuardDxe.inf
@@ -26,6 +26,9 @@
   Zydis/src/Utils.c
   Zydis/src/Zydis.c
 
+[Sources.X64]
+  X64/Cet.nasm
+
 [Packages]
   MdePkg/MdePkg.dec
   EfiGuardPkg/EfiGuardPkg.dec
diff --git a/EfiGuardDxe/EfiGuardDxe.vcxproj b/EfiGuardDxe/EfiGuardDxe.vcxproj
index 0b95d46..e6cc52c 100644
--- a/EfiGuardDxe/EfiGuardDxe.vcxproj
+++ b/EfiGuardDxe/EfiGuardDxe.vcxproj
@@ -23,6 +23,10 @@
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
   <Import Project="$(SolutionDir)\EfiGuard.props" />
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
+  <ImportGroup Label="PropertySheets">
+    <Import Project="$(VCTargetsPath)\BuildCustomizations\masm.props" />
+    <Import Project="$(VCTargetsPath)\BuildCustomizations\masm.targets" />
+  </ImportGroup>
   <ItemDefinitionGroup>
     <ClCompile>
       <PreprocessorDefinitions>ZYAN_NO_LIBC;ZYCORE_STATIC_BUILD;ZYDIS_STATIC_BUILD;ZYDIS_DISABLE_ENCODER;ZYDIS_DISABLE_FORMATTER;ZYDIS_DISABLE_AVX512;ZYDIS_DISABLE_KNC;%(PreprocessorDefinitions)</PreprocessorDefinitions>
@@ -73,6 +77,12 @@
     <ClCompile Include="Zydis\src\Zydis.c" />
   </ItemGroup>
   <ItemGroup>
+    <MASM Include="X64\Cet.asm" />
+  </ItemGroup>
+  <ItemGroup>
+    <None Include="X64\Cet.nasm" />
+  </ItemGroup>
+  <ItemGroup>
     <ClInclude Include="..\Include\Protocol\EfiGuard.h" />
     <ClInclude Include="arc.h" />
     <ClInclude Include="EfiGuardDxe.h" />
diff --git a/EfiGuardDxe/EfiGuardDxe.vcxproj.filters b/EfiGuardDxe/EfiGuardDxe.vcxproj.filters
index ef7eb16..e2bcab6 100644
--- a/EfiGuardDxe/EfiGuardDxe.vcxproj.filters
+++ b/EfiGuardDxe/EfiGuardDxe.vcxproj.filters
@@ -21,6 +21,9 @@
     <Filter Include="Header Files\Zydis\Internal">
       <UniqueIdentifier>{09843B9B-51DC-4418-9585-2ED4BD3F1643}</UniqueIdentifier>
     </Filter>
+    <Filter Include="Source Files\X64">
+      <UniqueIdentifier>{4cd0c7c8-71ff-4d6d-bb2d-a9e65cc3d7ce}</UniqueIdentifier>
+    </Filter>
     <Filter Include="Header Files\Protocol">
       <UniqueIdentifier>{aa6da080-fea5-447e-8722-35a98038eb4e}</UniqueIdentifier>
     </Filter>
@@ -97,6 +100,16 @@
     </ClCompile>
   </ItemGroup>
   <ItemGroup>
+    <MASM Include="X64\Cet.asm">
+      <Filter>Source Files\X64</Filter>
+    </MASM>
+  </ItemGroup>
+  <ItemGroup>
+    <None Include="X64\Cet.nasm">
+      <Filter>Source Files\X64</Filter>
+    </None>
+  </ItemGroup>
+  <ItemGroup>
     <ClInclude Include="EfiGuardDxe.h">
       <Filter>Header Files</Filter>
     </ClInclude>
diff --git a/EfiGuardDxe/X64/Cet.asm b/EfiGuardDxe/X64/Cet.asm
new file mode 100644
index 0000000..807fe6b
--- /dev/null
+++ b/EfiGuardDxe/X64/Cet.asm
@@ -0,0 +1,24 @@
+.code
+
+DisableCet PROC
+	; Pop pushed data for 'call'
+	mov rax, 1
+	incsspq rax
+
+	mov rax, cr4
+	btr eax, 23					; clear CR4_CET
+	mov cr4, rax
+	ret
+DisableCet ENDP
+
+EnableCet PROC
+	mov rax, cr4
+	bts eax, 23					; set CR4_CET
+	mov cr4, rax
+
+	; Use jmp to skip check for 'ret'
+	pop rax
+	jmp rax
+EnableCet ENDP
+
+end
diff --git a/EfiGuardDxe/X64/Cet.nasm b/EfiGuardDxe/X64/Cet.nasm
new file mode 100644
index 0000000..c31bc07
--- /dev/null
+++ b/EfiGuardDxe/X64/Cet.nasm
@@ -0,0 +1,23 @@
+DEFAULT REL
+SECTION .text
+
+global ASM_PFX(DisableCet)
+ASM_PFX(DisableCet):
+	; Pop pushed data for 'call'
+	mov rax, 1
+	incsspq rax
+
+	mov rax, cr4
+	btr eax, 23					; clear CR4_CET
+	mov cr4, rax
+	ret
+
+global ASM_PFX(EnableCet)
+ASM_PFX(EnableCet):
+	mov rax, cr4
+	bts eax, 23					; set CR4_CET
+	mov cr4, rax
+
+	; Use jmp to skip check for 'ret'
+	pop rax
+	jmp rax
diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index c15b7da..e7d88db 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -118,6 +118,40 @@ PrintKernelPatchInfo(
 	}
 }
 
+VOID
+EFIAPI
+DisableWriteProtect(
+	OUT BOOLEAN *WpEnabled,
+	OUT BOOLEAN *CetEnabled
+	)
+{
+	CONST UINTN Cr0 = AsmReadCr0();
+	*WpEnabled = (Cr0 & CR0_WP) != 0;
+	*CetEnabled = (AsmReadCr4() & CR4_CET) != 0;
+
+	if (*WpEnabled)
+	{
+		if (*CetEnabled)
+			DisableCet();
+		AsmWriteCr0(Cr0 & ~CR0_WP);
+	}
+}
+
+VOID
+EFIAPI
+EnableWriteProtect(
+	IN BOOLEAN WpEnabled,
+	IN BOOLEAN CetEnabled
+	)
+{
+	if (WpEnabled)
+	{
+		AsmWriteCr0(AsmReadCr0() | CR0_WP);
+		if (CetEnabled)
+			EnableCet();
+	}
+}
+
 VOID*
 EFIAPI
 CopyWpMem(
@@ -126,16 +160,12 @@ CopyWpMem(
 	IN UINTN Length
 	)
 {
-	CONST UINTN Cr0 = AsmReadCr0();
-	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
-	if (WpSet)
-		AsmWriteCr0(Cr0 & ~CR0_WP);
+	BOOLEAN WpEnabled, CetEnabled;
+	DisableWriteProtect(&WpEnabled, &CetEnabled);
 
 	VOID* Result = CopyMem(Destination, Source, Length);
 	
-	if (WpSet)
-		AsmWriteCr0(Cr0);
-
+	EnableWriteProtect(WpEnabled, CetEnabled);
 	return Result;
 }
 
@@ -147,16 +177,12 @@ SetWpMem(
 	IN UINT8 Value
 	)
 {
-	CONST UINTN Cr0 = AsmReadCr0();
-	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
-	if (WpSet)
-		AsmWriteCr0(Cr0 & ~CR0_WP);
+	BOOLEAN WpEnabled, CetEnabled;
+	DisableWriteProtect(&WpEnabled, &CetEnabled);
 
 	VOID* Result = SetMem(Destination, Length, Value);
 	
-	if (WpSet)
-		AsmWriteCr0(Cr0);
-
+	EnableWriteProtect(WpEnabled, CetEnabled);
 	return Result;
 }
 
diff --git a/EfiGuardDxe/util.h b/EfiGuardDxe/util.h
index b36692e..e1a049f 100644
--- a/EfiGuardDxe/util.h
+++ b/EfiGuardDxe/util.h
@@ -8,6 +8,7 @@
 
 #define CR0_WP			((UINTN)0x00010000) // CR0.WP
 #define CR0_PG			((UINTN)0x80000000) // CR0.PG
+#define CR4_CET			((UINTN)0x00800000) // CR4.CET
 #define CR4_LA57		((UINTN)0x00001000) // CR4.LA57
 #define MSR_EFER		((UINTN)0xC0000080) // Extended Function Enable Register
 #define EFER_LMA		((UINTN)0x00000400) // Long Mode Active
@@ -65,6 +66,45 @@ PrintKernelPatchInfo(
 	);
 
 //
+// Disables CET.
+//
+VOID
+EFIAPI
+DisableCet(
+	VOID
+	);
+
+//
+// Enables CET.
+//
+VOID
+EFIAPI
+EnableCet(
+	VOID
+	);
+
+//
+// Disables write protection if it is currently enabled.
+// Returns the current CET and WP states for use when calling EnableWriteProtect().
+//
+VOID
+EFIAPI
+DisableWriteProtect(
+	OUT BOOLEAN *WpEnabled,
+	OUT BOOLEAN *CetEnabled
+	);
+
+//
+// Enables write protection if it was previously enabled.
+//
+VOID
+EFIAPI
+EnableWriteProtect(
+	IN BOOLEAN WpEnabled,
+	IN BOOLEAN CetEnabled
+	);
+
+//
 // Wrapper for CopyMem() that disables write protection prior to copying if needed.
 //
 VOID*