Update vtoytool

[Ventoy.git] / DMPATCH / dmpatch.c

diff --git a/DMPATCH/dmpatch.c b/DMPATCH/dmpatch.c
index 359c5c3fb43b4ef2d3817af8caec474e1dea0a8c..e508a0624e976f9814054875c9489ced6db717fd 100644 (file)
--- a/DMPATCH/dmpatch.c
+++ b/DMPATCH/dmpatch.c
@@ -52,7 +52,8 @@ typedef struct ko_param
     unsigned long sym_put_addr;
     unsigned long sym_put_size;
     unsigned long kv_major;
-    unsigned long padding[2];
+    unsigned long ibt;
+    unsigned long padding[1];
 }ko_param;
 
 #pragma pack()
@@ -91,6 +92,52 @@ static volatile ko_param g_ko_param =
 #error "unsupported arch"
 #endif
 
+#ifdef VTOY_IBT
+#ifdef CONFIG_X86_64
+/* Using 64-bit values saves one instruction clearing the high half of low */
+#define DECLARE_ARGS(val, low, high)   unsigned long low, high
+#define EAX_EDX_VAL(val, low, high)    ((low) | (high) << 32)
+#define EAX_EDX_RET(val, low, high)    "=a" (low), "=d" (high)
+#else
+#define DECLARE_ARGS(val, low, high)   unsigned long long val
+#define EAX_EDX_VAL(val, low, high)    (val)
+#define EAX_EDX_RET(val, low, high)    "=A" (val)
+#endif
+
+#define        EX_TYPE_WRMSR                    8
+#define        EX_TYPE_RDMSR                    9
+#define MSR_IA32_S_CET                 0x000006a2 /* kernel mode cet */
+#define CET_ENDBR_EN                   (1ULL << 2)
+
+/* Exception table entry */
+#ifdef __ASSEMBLY__
+
+#define _ASM_EXTABLE_TYPE(from, to, type)                      \
+       .pushsection "__ex_table","a" ;                         \
+       .balign 4 ;                                             \
+       .long (from) - . ;                                      \
+       .long (to) - . ;                                        \
+       .long type ;                                            \
+       .popsection
+
+#else /* ! __ASSEMBLY__ */
+
+#define _ASM_EXTABLE_TYPE(from, to, type)                      \
+       " .pushsection \"__ex_table\",\"a\"\n"                  \
+       " .balign 4\n"                                          \
+       " .long (" #from ") - .\n"                              \
+       " .long (" #to ") - .\n"                                \
+       " .long " __stringify(type) " \n"                       \
+       " .popsection\n"
+
+#endif /* __ASSEMBLY__ */
+#endif /* VTOY_IBT */
+
+
+
+
+
+
 #define vdebug(fmt, args...) if(kprintf) kprintf(KERN_ERR fmt, ##args)
 
 static unsigned char *g_get_patch[MAX_PATCH] = { NULL };
@@ -166,11 +213,66 @@ static int notrace dmpatch_replace_code
     return 0;
 }
 
+#ifdef VTOY_IBT
+static __always_inline unsigned long long dmpatch_rdmsr(unsigned int msr)
+{
+       DECLARE_ARGS(val, low, high);
+
+       asm volatile("1: rdmsr\n"
+                    "2:\n"
+                    _ASM_EXTABLE_TYPE(1b, 2b, EX_TYPE_RDMSR)
+                    : EAX_EDX_RET(val, low, high) : "c" (msr));
+
+       return EAX_EDX_VAL(val, low, high);
+}
+
+static __always_inline void dmpatch_wrmsr(unsigned int msr, u32 low, u32 high)
+{
+       asm volatile("1: wrmsr\n"
+                    "2:\n"
+                    _ASM_EXTABLE_TYPE(1b, 2b, EX_TYPE_WRMSR)
+                    : : "c" (msr), "a"(low), "d" (high) : "memory");
+}
+
+static u64 dmpatch_ibt_save(void)
+{
+    u64 msr = 0;
+    u64 val = 0;
+
+    msr = dmpatch_rdmsr(MSR_IA32_S_CET);
+    val = msr & ~CET_ENDBR_EN;
+    dmpatch_wrmsr(MSR_IA32_S_CET, (u32)(val & 0xffffffffULL), (u32)(val >> 32));
+
+    return msr;
+}
+
+static void dmpatch_ibt_restore(u64 save)
+{
+       u64 msr;
+
+    msr = dmpatch_rdmsr(MSR_IA32_S_CET);
+
+       msr &= ~CET_ENDBR_EN;
+       msr |= (save & CET_ENDBR_EN);
+
+    dmpatch_wrmsr(MSR_IA32_S_CET, (u32)(msr & 0xffffffffULL), (u32)(msr >> 32));
+}
+#else
+static u64 dmpatch_ibt_save(void) { return 0; }
+static void dmpatch_ibt_restore(u64 save) { (void)save; }
+#endif
+
 static int notrace dmpatch_init(void)
 {
     int r = 0;
     int rc = 0;
-
+    u64 msr = 0;
+    
+    if (g_ko_param.ibt == 0x8888)
+    {
+        msr = dmpatch_ibt_save();
+    }
+    
     kprintf = (printk_pf)(g_ko_param.printk_addr); 
 
     vdebug("dmpatch_init start pagesize=%lu ...\n", g_ko_param.pgsize);
@@ -218,6 +320,11 @@ static int notrace dmpatch_init(void)
     vdebug("######## dm patch success ###########\n");
     vdebug("#####################################\n");
 
+    if (g_ko_param.ibt == 0x8888)
+    {
+        dmpatch_ibt_restore(msr);
+    }
+
 out:
 
        return rc;
@@ -226,6 +333,12 @@ out:
 static void notrace dmpatch_exit(void)
 {
     int i = 0;
+    u64 msr;
+
+    if (g_ko_param.ibt == 0x8888)
+    {
+        msr = dmpatch_ibt_save();
+    }
 
     for (i = 0; i < MAX_PATCH; i++)
     {
@@ -234,6 +347,11 @@ static void notrace dmpatch_exit(void)
     }
 
     vdebug("dmpatch_exit success\n");
+
+    if (g_ko_param.ibt == 0x8888)
+    {
+        dmpatch_ibt_restore(msr);
+    }
 }
 
 module_init(dmpatch_init);