From ff33090fd09dc9a0617a39ce2cee2ec9028e985d Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Mon, 15 Feb 2021 16:48:34 +0100
Subject: [PATCH] Check whether RCX is canonical in sysretq.

---
 src/arch/x86_64/interrupt/syscall.rs | 73 ++++++++++++++++++++++++----
 1 file changed, 63 insertions(+), 10 deletions(-)

diff --git a/src/arch/x86_64/interrupt/syscall.rs b/src/arch/x86_64/interrupt/syscall.rs
index ef8ec59..a404728 100644
--- a/src/arch/x86_64/interrupt/syscall.rs
+++ b/src/arch/x86_64/interrupt/syscall.rs
@@ -58,16 +58,29 @@ pub unsafe extern "C" fn __inner_syscall_instruction(stack: *mut InterruptStack)
     });
 }
 
+macro_rules! fast_sysretq {
+    () => { "
+        pop rcx                 // Pop userspace return pointer
+        sub rsp, 8              // Pop fake userspace CS
+        pop r11                 // Pop rflags
+        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
+        mov rsp, gs:[0x70]      // Restore userspace stack pointer
+        swapgs                  // Restore gs from TSS to user data
+        ud2
+        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
+    " }
+}
+
 function!(syscall_instruction => {
     // Yes, this is magic. No, you don't need to understand
     "
         swapgs                    // Set gs segment to TSS
         mov gs:[0x70], rsp        // Save userspace stack pointer
         mov rsp, gs:[4]           // Load kernel stack pointer
-        push QWORD PTR 5 * 8 + 3  // Push fake userspace SS (resembling iret frame)
+        push QWORD PTR 5 * 8 + 3  // Push fake SS (resembling iret stack frame)
         push QWORD PTR gs:[0x70]  // Push userspace rsp
         push r11                  // Push rflags
-        push QWORD PTR 6 * 8 + 3  // Push fake userspace CS (resembling iret frame)
+        push QWORD PTR 6 * 8 + 3  // Push fake CS (resembling iret stack frame)
         push rcx                  // Push userspace return pointer
     ",
 
@@ -94,14 +107,54 @@ function!(syscall_instruction => {
 
     // Return
     "
-        pop rcx                 // Pop userspace return pointer
-        add rsp, 8              // Pop CS
-        pop r11                 // Pop rflags
-        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
-        add rsp, 8              // Pop SS
-        mov rsp, gs:[0x70]      // Restore userspace stack pointer
-        swapgs                  // Restore gs from TSS to user data
-        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
+        // Test whether RCX is canonical. This is not strictly necessary, but could be fatal if
+        // some kernel bug would allow RCX to be modified by user code.
+
+        // Peek at the preserved RCX, to double-check that it is canonical.
+        // Set ZF iff bit 47 (i.e. the bit that must be sign extended) is set.
+
+        // TODO: Any hacky one-instruction canonicalness check?
+
+        bt QWORD PTR [rsp], 47
+
+        // Jump to the canonicalness check.
+        jnz 2f
+
+        // Otherwise, continue with the fast return.
+        1:
+        // Fast sysretq
+    ",
+
+    fast_sysretq!(),
+
+    "
+        // Jumps to aligned addresses are faster, and it is not always the case the userspace
+        // doesn't call from higher half (though unlikely).
+        .align 8, 0x90
+
+        2:
+
+        // Load userspace return RIP.
+        mov rcx, [rsp]
+
+        // Shift it to the right by 47, only getting the bits which must all be zero or one.
+        shr rcx, 47
+        // Since we have already checked that RCX[47] == 1, RCX[63:47] must all be one.
+        cmp rcx, 0x1FFFF
+        // If the address is canonical, go to the slower iretq.
+        jne 3f
+    ",
+
+        fast_sysretq!(),
+
+    "
+        3:
+        // Slow iretq
+
+        xor rcx, rcx
+        xor r11, r11
+        swapgs
+        iretq
     ",
 });