From 3d2e032b41c27c10f0aa719d2731741df037840e Mon Sep 17 00:00:00 2001
From: Moritz Lipp <mlipp@amazon.com>
Date: Tue, 28 Jun 2022 19:56:09 +0200
Subject: [PATCH 1/2] Initial support for kernels with IBT

---
 module/pteditor.c | 54 ++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/module/pteditor.c b/module/pteditor.c
index 8637d7f..bc286eb 100644
--- a/module/pteditor.c
+++ b/module/pteditor.c
@@ -13,6 +13,33 @@
 #include <linux/proc_fs.h>
 #include <linux/kprobes.h>
 
+#ifdef CONFIG_X86_KERNEL_IBT
+__noendbr u64 ibt_save(void)
+{
+	u64 msr = 0;
+
+	if (cpu_feature_enabled(X86_FEATURE_IBT)) {
+		rdmsrl(MSR_IA32_S_CET, msr);
+		wrmsrl(MSR_IA32_S_CET, msr & ~CET_ENDBR_EN);
+	}
+
+	return msr;
+}
+
+__noendbr void ibt_restore(u64 save)
+{
+	u64 msr;
+
+	if (cpu_feature_enabled(X86_FEATURE_IBT)) {
+		rdmsrl(MSR_IA32_S_CET, msr);
+		msr &= ~CET_ENDBR_EN;
+		msr |= (save & CET_ENDBR_EN);
+		wrmsrl(MSR_IA32_S_CET, msr);
+	}
+}
+
+#endif
+
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 8, 0)
 #include <linux/mmap_lock.h>
 #endif
@@ -105,7 +132,23 @@ static inline int pmd_large(pmd_t pmd) {
 #define KPROBE_KALLSYMS_LOOKUP 1
 typedef unsigned long (*kallsyms_lookup_name_t)(const char *name);
 kallsyms_lookup_name_t kallsyms_lookup_name_func;
+
+#ifdef CONFIG_X86_KERNEL_IBT
+unsigned long kallsyms_lookup_name_func_ibt(const char* name) {
+  u64 ibt;
+  unsigned long r;
+
+  ibt = ibt_save();
+  r = kallsyms_lookup_name_func(name);
+  ibt_restore(ibt);
+
+  return r;
+}
+
+#define kallsyms_lookup_name kallsyms_lookup_name_func_ibt
+#else
 #define kallsyms_lookup_name kallsyms_lookup_name_func
+#endif
 
 static struct kprobe kp = {
     .symbol_name = "kallsyms_lookup_name"
@@ -679,11 +722,16 @@ static int __init pteditor_init(void) {
 #endif
 
 #ifdef KPROBE_KALLSYMS_LOOKUP
-    register_kprobe(&kp);
-    kallsyms_lookup_name = (kallsyms_lookup_name_t) kp.addr;
+    r = register_kprobe(&kp);
+    if (r != 0) {
+      pr_alert("Could not register kprobe\n");
+      return -ENXIO;
+    }
+
+    kallsyms_lookup_name_func = (kallsyms_lookup_name_t) kp.addr;
     unregister_kprobe(&kp);
 
-    if(!unlikely(kallsyms_lookup_name)) {
+    if(!unlikely(kallsyms_lookup_name_func)) {
       pr_alert("Could not retrieve kallsyms_lookup_name address\n");
       return -ENXIO;
     }

From 5f3bf1071c3507e039e23f63948cf6c81bf60ec4 Mon Sep 17 00:00:00 2001
From: Moritz Lipp <mlipp@amazon.com>
Date: Wed, 29 Jun 2022 09:45:03 +0200
Subject: [PATCH 2/2] Handle IBT for invalidate_tlb_kernel

---
 module/pteditor.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/module/pteditor.c b/module/pteditor.c
index bc286eb..db0136a 100644
--- a/module/pteditor.c
+++ b/module/pteditor.c
@@ -257,7 +257,14 @@ void _flush_tlb_page_smp(void* info) {
 static void
 invalidate_tlb_kernel(unsigned long addr) {
 #if defined(__i386__) || defined(__x86_64__)
+#ifdef CONFIG_X86_KERNEL_IBT
+  u64 ibt;
+  ibt = ibt_save();
+#endif
   flush_tlb_mm_range_func(get_mm(task_pid_nr(current)), addr, addr + real_page_size, real_page_shift, false);
+#ifdef CONFIG_X86_KERNEL_IBT
+  ibt_restore(ibt);
+#endif
 #elif defined(__aarch64__)
   struct vm_area_struct *vma = find_vma(current->mm, addr);
   tlb_page_t tlb_page;