[kvm-unit-tests PATCH 10/14] x86: Look up the PTEs rather than assuming them

[Aaron Lewis <aaronlewis@xxxxxxxxxx>]

Rather than assuming which PTEs the SMEP test runs on, look them up to
ensure they are correct.  If this test were to run on a different page
table (ie: run in an L2 test) the wrong PTEs would be set.  Switch to
looking up the PTEs to avoid this from happening.

Signed-off-by: Aaron Lewis <aaronlewis@xxxxxxxxxx>
---
 lib/libcflat.h |  1 +
 lib/x86/vm.c   | 21 +++++++++++++++++++++
 lib/x86/vm.h   |  3 +++
 x86/access.c   | 26 ++++++++++++++++++--------
 x86/cstart64.S |  1 -
 x86/flat.lds   |  1 +
 6 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/lib/libcflat.h b/lib/libcflat.h
index 9bb7e08..c1fd31f 100644
--- a/lib/libcflat.h
+++ b/lib/libcflat.h
@@ -35,6 +35,7 @@
 #define __ALIGN_MASK(x, mask)	(((x) + (mask)) & ~(mask))
 #define __ALIGN(x, a)		__ALIGN_MASK(x, (typeof(x))(a) - 1)
 #define ALIGN(x, a)		__ALIGN((x), (a))
+#define ALIGN_DOWN(x, a)	__ALIGN((x) - ((a) - 1), (a))
 #define IS_ALIGNED(x, a)	(((x) & ((typeof(x))(a) - 1)) == 0)
 
 #define MIN(a, b)		((a) < (b) ? (a) : (b))
diff --git a/lib/x86/vm.c b/lib/x86/vm.c
index 5cd2ee4..6a70ef6 100644
--- a/lib/x86/vm.c
+++ b/lib/x86/vm.c
@@ -281,3 +281,24 @@ void force_4k_page(void *addr)
 	if (pte & PT_PAGE_SIZE_MASK)
 		split_large_page(ptep, 2);
 }
+
+/*
+ * Call the callback on each page from virt to virt + len.
+ */
+void walk_pte(void *virt, size_t len, pte_callback_t callback)
+{
+    pgd_t *cr3 = current_page_table();
+    uintptr_t start = (uintptr_t)virt;
+    uintptr_t end = (uintptr_t)virt + len;
+    struct pte_search search;
+    size_t page_size;
+    uintptr_t curr;
+
+    for (curr = start; curr < end; curr = ALIGN_DOWN(curr + page_size, page_size)) {
+        search = find_pte_level(cr3, (void *)curr, 1);
+        assert(found_leaf_pte(search));
+        page_size = 1ul << PGDIR_BITS(search.level);
+
+        callback(search, (void *)curr);
+    }
+}
diff --git a/lib/x86/vm.h b/lib/x86/vm.h
index d9753c3..4c6dff9 100644
--- a/lib/x86/vm.h
+++ b/lib/x86/vm.h
@@ -52,4 +52,7 @@ struct vm_vcpu_info {
         u64 cr0;
 };
 
+typedef void (*pte_callback_t)(struct pte_search search, void *va);
+void walk_pte(void *virt, size_t len, pte_callback_t callback);
+
 #endif
diff --git a/x86/access.c b/x86/access.c
index a781a0c..8e3a718 100644
--- a/x86/access.c
+++ b/x86/access.c
@@ -201,10 +201,24 @@ static void set_cr0_wp(int wp)
     }
 }
 
+static void clear_user_mask(struct pte_search search, void *va)
+{
+	*search.pte &= ~PT_USER_MASK;
+}
+
+static void set_user_mask(struct pte_search search, void *va)
+{
+	*search.pte |= PT_USER_MASK;
+
+	/* Flush to avoid spurious #PF */
+	invlpg(va);
+}
+
 static unsigned set_cr4_smep(int smep)
 {
+    extern char stext, etext;
+    size_t len = (size_t)&etext - (size_t)&stext;
     unsigned long cr4 = shadow_cr4;
-    extern u64 ptl2[];
     unsigned r;
 
     cr4 &= ~CR4_SMEP_MASK;
@@ -214,14 +228,10 @@ static unsigned set_cr4_smep(int smep)
         return 0;
 
     if (smep)
-        ptl2[2] &= ~PT_USER_MASK;
+        walk_pte(&stext, len, clear_user_mask);
     r = write_cr4_checking(cr4);
-    if (r || !smep) {
-        ptl2[2] |= PT_USER_MASK;
-
-	/* Flush to avoid spurious #PF */
-	invlpg((void *)(2 << 21));
-    }
+    if (r || !smep)
+        walk_pte(&stext, len, set_user_mask);
     if (!r)
         shadow_cr4 = cr4;
     return r;
diff --git a/x86/cstart64.S b/x86/cstart64.S
index ddb83a0..ff79ae7 100644
--- a/x86/cstart64.S
+++ b/x86/cstart64.S
@@ -17,7 +17,6 @@ stacktop:
 .data
 
 .align 4096
-.globl ptl2
 ptl2:
 i = 0
 	.rept 512 * 4
diff --git a/x86/flat.lds b/x86/flat.lds
index a278b56..337bc44 100644
--- a/x86/flat.lds
+++ b/x86/flat.lds
@@ -3,6 +3,7 @@ SECTIONS
     . = 4M + SIZEOF_HEADERS;
     stext = .;
     .text : { *(.init) *(.text) *(.text.*) }
+    etext = .;
     . = ALIGN(4K);
     .data : {
           *(.data)
-- 
2.34.0.rc1.387.gb447b232ab-goog