Merge pull request #41 from d-we/master

Fix TLB invalidation for PID != 0

README.md

@@ -139,7 +139,8 @@ System Info | Descriptions
 
  TLB/Barriers       | Descriptions
 --------------------------------|---------------------------------------------
-`void `[`ptedit_invalidate_tlb`](#group__BARRIERS_1gad2d64fa589bc626ba41ccf18c60d159f)`(void * address)`            | Invalidates the TLB for a given address on all CPUs.
+`void `[`ptedit_invalidate_tlb`](#group__BARRIERS_1gad2d64fa589bc626ba41ccf18c60d159f)`(void * address)`            | Invalidates the TLB entry of current process for a given address on all CPUs.
+`void `[`ptedit_invalidate_tlb_pid`](#group__BARRIERS_1gad2d64fa589bc626ba41ccf18c60d159f)`(pid_t pid, void * address)`            | Invalidates the TLB for a given PID and address on all CPUs.
 `void `[`ptedit_full_serializing_barrier`](#group__BARRIERS_1ga35efff6b34856596b467ef3a5075adc6)`()`            | A full serializing barrier which stops everything.
 
  Memory types (PATs/MAIRs)       | Descriptions

module/pteditor.c

@@ -134,7 +134,7 @@ typedef struct {
 static bool device_busy = false;
 static bool mm_is_locked = false;
 
-void (*invalidate_tlb)(unsigned long);
+void (*invalidate_tlb)(pid_t, void*);
 void (*flush_tlb_mm_range_func)(struct mm_struct*, unsigned long, unsigned long, unsigned int, bool);
 void (*native_write_cr4_func)(unsigned long);
 static struct mm_struct* get_mm(size_t);
@@ -203,8 +203,8 @@ _invalidate_tlb(void *addr) {
 }
 
 static void
-invalidate_tlb_custom(unsigned long addr) {
-  on_each_cpu(_invalidate_tlb, (void*) addr, 1);
+invalidate_tlb_custom(pid_t pid, void* addr) {
+  on_each_cpu(_invalidate_tlb, addr, 1);
 }
 
 #if defined(__aarch64__)
@@ -220,9 +220,9 @@ void _flush_tlb_page_smp(void* info) {
 #endif
 
 static void
-invalidate_tlb_kernel(unsigned long addr) {
+invalidate_tlb_kernel(pid_t pid, void* addr) {
 #if defined(__i386__) || defined(__x86_64__)
-  flush_tlb_mm_range_func(get_mm(task_pid_nr(current)), addr, addr + real_page_size, real_page_shift, false);
+  flush_tlb_mm_range_func(get_mm(pid), (unsigned long) addr, (unsigned long) addr + real_page_size, real_page_shift, false);
 #elif defined(__aarch64__)
   struct vm_area_struct *vma = find_vma(current->mm, addr);
   tlb_page_t tlb_page;
@@ -421,7 +421,7 @@ static int update_vm(ptedit_entry_t* new_entry, int lock) {
       set_pte(old_entry.pte, native_make_pte(new_entry->pte));
   }
 
-  invalidate_tlb(addr);
+  invalidate_tlb(old_entry.pid, (void*) addr);
 
   /* Unlock mm */
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 8, 0)
@@ -585,9 +585,20 @@ static long device_ioctl(struct file *file, unsigned int ioctl_num, unsigned lon
     }
     case PTEDITOR_IOCTL_CMD_GET_PAGESIZE:
         return real_page_size;
+    case PTEDITOR_IOCTL_CMD_INVALIDATE_TLB_PID:
+    {
+        ptedit_invalidate_tlb_args_t args;
+        (void)from_user(&args, (void*)ioctl_param, sizeof(args));
+        invalidate_tlb(args.pid, args.address);
+        return 0;
+    }
     case PTEDITOR_IOCTL_CMD_INVALIDATE_TLB:
-        invalidate_tlb(ioctl_param);
+    {
+        // this is implemented as its own call to stay backwards compatible
+        // even in case a user uses the old ioctl calls
+        invalidate_tlb(task_pid_nr(current), (void*) ioctl_param);
         return 0;
+    }
     case PTEDITOR_IOCTL_CMD_GET_PAT:
     {
 #if defined(__i386__) || defined(__x86_64__)

module/pteditor.h

@@ -90,6 +90,14 @@ typedef struct {
     size_t root;
 } ptedit_paging_t;
 
+/**
+ * Structure to hold the arguments for TLB invalidation
+ */
+typedef struct {
+    pid_t pid;
+    void* address;
+} ptedit_invalidate_tlb_args_t;
+
 #define PTEDIT_VALID_MASK_PGD (1<<0)
 #define PTEDIT_VALID_MASK_P4D (1<<1)
 #define PTEDIT_VALID_MASK_PUD (1<<2)
@@ -140,6 +148,9 @@ typedef struct {
 
 #define PTEDITOR_IOCTL_CMD_SWITCH_TLB_INVALIDATION \
   _IOR(PTEDITOR_IOCTL_MAGIC_NUMBER, 13, size_t)
+
+#define PTEDITOR_IOCTL_CMD_INVALIDATE_TLB_PID \
+  _IOR(PTEDITOR_IOCTL_MAGIC_NUMBER, 14, size_t)
 #else
 #define PTEDITOR_READ_PAGE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define PTEDITOR_WRITE_PAGE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_BUFFERED, FILE_READ_DATA)

ptedit.c

@@ -275,20 +275,20 @@ ptedit_fnc void ptedit_update_user_ext(void* address, pid_t pid, ptedit_entry_t*
         pset(root + pgdi * ptedit_entry_size, vm->pgd);
     }
 
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 // ---------------------------------------------------------------------------
 static void ptedit_update_user(void* address, pid_t pid, ptedit_entry_t* vm) {
     ptedit_update_user_ext(address, pid, vm, ptedit_phys_write_pwrite);
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 
 // ---------------------------------------------------------------------------
 static void ptedit_update_user_map(void* address, pid_t pid, ptedit_entry_t* vm) {
     ptedit_update_user_ext(address, pid, vm, ptedit_phys_write_map);
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 // ---------------------------------------------------------------------------
@@ -642,9 +642,24 @@ void ptedit_set_paging_root(pid_t pid, size_t root) {
 #endif
 }
 
+// ---------------------------------------------------------------------------
+ptedit_fnc void ptedit_invalidate_tlb_pid(pid_t pid, void* address) {
+#if defined(LINUX)
+    ptedit_invalidate_tlb_args_t args;
+    args.pid = pid;
+    args.address = address;
+    ioctl(ptedit_fd, PTEDITOR_IOCTL_CMD_INVALIDATE_TLB_PID, (size_t)&args, pid);
+#else
+    size_t vaddr = (size_t)address;
+    DWORD returnLength;
+    DeviceIoControl(ptedit_fd, PTEDITOR_FLUSH_TLB, (LPVOID)&vaddr, sizeof(vaddr), (LPVOID)&vaddr, sizeof(vaddr), &returnLength, 0);
+#endif
+}
 
 // ---------------------------------------------------------------------------
 ptedit_fnc void ptedit_invalidate_tlb(void* address) {
+    // we do not directly call ptedit_invalidate_tlb_pid to ensure that the old
+    // API is still working (for backwards compatibility)
 #if defined(LINUX)
     ioctl(ptedit_fd, PTEDITOR_IOCTL_CMD_INVALIDATE_TLB, (size_t)address);
 #else

ptedit.h

@@ -632,13 +632,21 @@ ptedit_fnc void ptedit_set_paging_root(pid_t pid, size_t root);
  */
 
  /**
-  * Invalidates the TLB for a given address on all CPUs.
+  * Invalidates the TLB for a given address (belonging to the current process) on all CPUs.
   *
   * @param[in] address The address to invalidate
   *
   */
 ptedit_fnc void ptedit_invalidate_tlb(void* address);
 
+ /**
+  * Invalidates the TLB for a given address (belonging to the specified pid) on all CPUs.
+  *
+  * @param[in] address The address to invalidate
+  *
+  */
+ptedit_fnc void ptedit_invalidate_tlb_pid(pid_t pid, void* address);
+
  /**
   * Change the method used for flushing the TLB (either kernel or custom function)
   *

ptedit_header.h

@@ -92,6 +92,14 @@ typedef struct {
     size_t root;
 } ptedit_paging_t;
 
+/**
+ * Structure to hold the arguments for TLB invalidation
+ */
+typedef struct {
+    pid_t pid;
+    void* address;
+} ptedit_invalidate_tlb_args_t;
+
 #define PTEDIT_VALID_MASK_PGD (1<<0)
 #define PTEDIT_VALID_MASK_P4D (1<<1)
 #define PTEDIT_VALID_MASK_PUD (1<<2)
@@ -142,6 +150,9 @@ typedef struct {
 
 #define PTEDITOR_IOCTL_CMD_SWITCH_TLB_INVALIDATION \
   _IOR(PTEDITOR_IOCTL_MAGIC_NUMBER, 13, size_t)
+
+#define PTEDITOR_IOCTL_CMD_INVALIDATE_TLB_PID \
+  _IOR(PTEDITOR_IOCTL_MAGIC_NUMBER, 14, size_t)
 #else
 #define PTEDITOR_READ_PAGE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define PTEDITOR_WRITE_PAGE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_BUFFERED, FILE_READ_DATA)
@@ -789,13 +800,21 @@ ptedit_fnc void ptedit_set_paging_root(pid_t pid, size_t root);
  */
 
  /**
-  * Invalidates the TLB for a given address on all CPUs.
+  * Invalidates the TLB for a given address (belonging to the current process) on all CPUs.
   *
   * @param[in] address The address to invalidate
   *
   */
 ptedit_fnc void ptedit_invalidate_tlb(void* address);
 
+ /**
+  * Invalidates the TLB for a given address (belonging to the specified pid) on all CPUs.
+  *
+  * @param[in] address The address to invalidate
+  *
+  */
+ptedit_fnc void ptedit_invalidate_tlb_pid(pid_t pid, void* address);
+
  /**
   * Change the method used for flushing the TLB (either kernel or custom function)
   *
@@ -1255,20 +1274,20 @@ ptedit_fnc void ptedit_update_user_ext(void* address, pid_t pid, ptedit_entry_t*
         pset(root + pgdi * ptedit_entry_size, vm->pgd);
     }
 
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 // ---------------------------------------------------------------------------
 static void ptedit_update_user(void* address, pid_t pid, ptedit_entry_t* vm) {
     ptedit_update_user_ext(address, pid, vm, ptedit_phys_write_pwrite);
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 
 // ---------------------------------------------------------------------------
 static void ptedit_update_user_map(void* address, pid_t pid, ptedit_entry_t* vm) {
     ptedit_update_user_ext(address, pid, vm, ptedit_phys_write_map);
-    ptedit_invalidate_tlb(address);
+    ptedit_invalidate_tlb_pid(pid, address);
 }
 
 // ---------------------------------------------------------------------------
@@ -1622,9 +1641,24 @@ void ptedit_set_paging_root(pid_t pid, size_t root) {
 #endif
 }
 
+// ---------------------------------------------------------------------------
+ptedit_fnc void ptedit_invalidate_tlb_pid(pid_t pid, void* address) {
+#if defined(LINUX)
+    ptedit_invalidate_tlb_args_t args;
+    args.pid = pid;
+    args.address = address;
+    ioctl(ptedit_fd, PTEDITOR_IOCTL_CMD_INVALIDATE_TLB_PID, (size_t)&args, pid);
+#else
+    size_t vaddr = (size_t)address;
+    DWORD returnLength;
+    DeviceIoControl(ptedit_fd, PTEDITOR_FLUSH_TLB, (LPVOID)&vaddr, sizeof(vaddr), (LPVOID)&vaddr, sizeof(vaddr), &returnLength, 0);
+#endif
+}
 
 // ---------------------------------------------------------------------------
 ptedit_fnc void ptedit_invalidate_tlb(void* address) {
+    // we do not directly call ptedit_invalidate_tlb_pid to ensure that the old
+    // API is still working (for backwards compatibility)
 #if defined(LINUX)
     ioctl(ptedit_fd, PTEDITOR_IOCTL_CMD_INVALIDATE_TLB, (size_t)address);
 #else