From 2c82592ab463f1f38237919a12145f34eaadda23 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@microsoft.com>
Date: Tue, 19 Mar 2024 08:47:28 -0700
Subject: [PATCH] GH-116017: Put JIT code and data on the same page (GH-116845)

---
 Python/jit.c           | 52 +++++++++++++++++++-----------------------
 Tools/jit/_stencils.py |  7 +++---
 Tools/jit/_targets.py  |  2 +-
 3 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/Python/jit.c b/Python/jit.c
index dae25166b1f..f67d641fe12 100644
--- a/Python/jit.c
+++ b/Python/jit.c
@@ -112,26 +112,6 @@ mark_executable(unsigned char *memory, size_t size)
     return 0;
 }
 
-static int
-mark_readable(unsigned char *memory, size_t size)
-{
-    if (size == 0) {
-        return 0;
-    }
-    assert(size % get_page_size() == 0);
-#ifdef MS_WINDOWS
-    DWORD old;
-    int failed = !VirtualProtect(memory, size, PAGE_READONLY, &old);
-#else
-    int failed = mprotect(memory, size, PROT_READ);
-#endif
-    if (failed) {
-        jit_error("unable to protect readable memory");
-        return -1;
-    }
-    return 0;
-}
-
 // JIT compiler stuff: /////////////////////////////////////////////////////////
 
 // Warning! AArch64 requires you to get your hands dirty. These are your gloves:
@@ -409,12 +389,14 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size
         code_size += group->code.body_size;
         data_size += group->data.body_size;
     }
-    // Round up to the nearest page (code and data need separate pages):
+    code_size += stencil_groups[_FATAL_ERROR].code.body_size;
+    data_size += stencil_groups[_FATAL_ERROR].data.body_size;
+    // Round up to the nearest page:
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
-    code_size += page_size - (code_size & (page_size - 1));
-    data_size += page_size - (data_size & (page_size - 1));
-    unsigned char *memory = jit_alloc(code_size + data_size);
+    size_t padding = page_size - ((code_size + data_size) & (page_size - 1));
+    size_t total_size = code_size + data_size + padding;
+    unsigned char *memory = jit_alloc(total_size);
     if (memory == NULL) {
         return -1;
     }
@@ -444,14 +426,26 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size
         code += group->code.body_size;
         data += group->data.body_size;
     }
-    if (mark_executable(memory, code_size) ||
-        mark_readable(memory + code_size, data_size))
-    {
-        jit_free(memory, code_size + data_size);
+    // Protect against accidental buffer overrun into data:
+    const StencilGroup *group = &stencil_groups[_FATAL_ERROR];
+    uint64_t patches[] = GET_PATCHES();
+    patches[HoleValue_CODE] = (uint64_t)code;
+    patches[HoleValue_CONTINUE] = (uint64_t)code;
+    patches[HoleValue_DATA] = (uint64_t)data;
+    patches[HoleValue_EXECUTOR] = (uint64_t)executor;
+    patches[HoleValue_TOP] = (uint64_t)code;
+    patches[HoleValue_ZERO] = 0;
+    emit(group, patches);
+    code += group->code.body_size;
+    data += group->data.body_size;
+    assert(code == memory + code_size);
+    assert(data == memory + code_size + data_size);
+    if (mark_executable(memory, total_size)) {
+        jit_free(memory, total_size);
         return -1;
     }
     executor->jit_code = memory;
-    executor->jit_size = code_size + data_size;
+    executor->jit_size = total_size;
     return 0;
 }
 
diff --git a/Tools/jit/_stencils.py b/Tools/jit/_stencils.py
index 78c566d9c8a..05c4ce8249f 100644
--- a/Tools/jit/_stencils.py
+++ b/Tools/jit/_stencils.py
@@ -124,7 +124,7 @@ class Stencil:
         ):
             self.holes.append(hole.replace(offset=base + 4 * i, kind=kind))
 
-    def remove_jump(self) -> None:
+    def remove_jump(self, *, alignment: int = 1) -> None:
         """Remove a zero-length continuation jump, if it exists."""
         hole = max(self.holes, key=lambda hole: hole.offset)
         match hole:
@@ -170,7 +170,7 @@ class Stencil:
                 offset -= 2
             case _:
                 return
-        if self.body[offset:] == jump:
+        if self.body[offset:] == jump and offset % alignment == 0:
             self.body = self.body[:offset]
             self.holes.remove(hole)
 
@@ -199,9 +199,8 @@ class StencilGroup:
             ):
                 self.code.pad(alignment)
                 self.code.emit_aarch64_trampoline(hole)
-                self.code.pad(alignment)
                 self.code.holes.remove(hole)
-        self.code.remove_jump()
+        self.code.remove_jump(alignment=alignment)
         self.code.pad(alignment)
         self.data.pad(8)
         for stencil in [self.code, self.data]:
diff --git a/Tools/jit/_targets.py b/Tools/jit/_targets.py
index 417fdb56ccf..66db3586792 100644
--- a/Tools/jit/_targets.py
+++ b/Tools/jit/_targets.py
@@ -89,7 +89,7 @@ class _Target(typing.Generic[_S, _R]):
         if group.data.body:
             line = f"0: {str(bytes(group.data.body)).removeprefix('b')}"
             group.data.disassembly.append(line)
-        group.process_relocations()
+        group.process_relocations(alignment=self.alignment)
         return group
 
     def _handle_section(self, section: _S, group: _stencils.StencilGroup) -> None: