Nmi branch test

arch/x86/include/asm/asm-prototypes.h

@@ -6,6 +6,7 @@
 #include <asm/page.h>
 #include <asm/checksum.h>
 #include <asm/mce.h>
+#include <asm/rex.h>
 
 #include <asm-generic/asm-prototypes.h>
 

arch/x86/include/asm/rex.h

@@ -0,0 +1,40 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef _ASM_X86_REX_UNWIND_H
+#define _ASM_X86_REX_UNWIND_H
+
+#include <linux/linkage.h>
+#include <linux/percpu.h>
+
+#ifndef __ASSEMBLY__
+
+/*
+ * Total stack is 8 pages (32k) large. 4 pages are reserved for kernel helpers/
+ * Therefore, the actual usable stack is 4 pages.
+ */
+#define REX_STACK_ORDER 3
+#define REX_STACK_SIZE (PAGE_SIZE << REX_STACK_ORDER)
+
+struct bpf_prog;
+struct bpf_insn;
+
+DECLARE_PER_CPU(unsigned char, rex_termination_state);
+DECLARE_PER_CPU(void *, rex_stack_ptr);
+
+extern asmlinkage unsigned int rex_dispatcher_func(
+	const void *ctx,
+	const struct bpf_prog *prog,
+	unsigned int (*bpf_func)(const void *,
+				 const struct bpf_insn *));
+
+int arch_init_rex_stack(void);
+
+static __always_inline bool arch_on_rex_stack(struct pt_regs *regs)
+{
+	unsigned long sp = regs->sp;
+	u64 rex_tos = (u64)this_cpu_read_stable(rex_stack_ptr);
+	return sp >= (rex_tos + 8 - REX_STACK_SIZE) && sp < (rex_tos + 8);
+}
+
+#endif /* !__ASSEMBLY__ */
+
+#endif /* _ASM_X86_REX_UNWIND_H */

arch/x86/include/asm/stacktrace.h

@@ -19,6 +19,7 @@ enum stack_type {
 	STACK_TYPE_IRQ,
 	STACK_TYPE_SOFTIRQ,
 	STACK_TYPE_ENTRY,
+	STACK_TYPE_REX,
 	STACK_TYPE_EXCEPTION,
 	STACK_TYPE_EXCEPTION_LAST = STACK_TYPE_EXCEPTION + N_EXCEPTION_STACKS-1,
 };

arch/x86/kernel/dumpstack_64.c

@@ -17,6 +17,7 @@
 #include <linux/nmi.h>
 
 #include <asm/cpu_entry_area.h>
+#include <asm/rex.h>
 #include <asm/stacktrace.h>
 
 static const char * const exception_stack_names[] = {
@@ -50,6 +51,9 @@ const char *stack_type_name(enum stack_type type)
 		return "ENTRY_TRAMPOLINE";
 	}
 
+	if (type == STACK_TYPE_REX)
+		return "REX";
+
 	if (type >= STACK_TYPE_EXCEPTION && type <= STACK_TYPE_EXCEPTION_LAST)
 		return exception_stack_names[type - STACK_TYPE_EXCEPTION];
 
@@ -167,6 +171,41 @@ static __always_inline bool in_irq_stack(unsigned long *stack, struct stack_info
 	return true;
 }
 
+static __always_inline bool in_rex_stack(unsigned long *stack, struct stack_info *info)
+{
+	unsigned long *end = (unsigned long *)this_cpu_read(rex_stack_ptr);
+	unsigned long *begin;
+
+	/*
+	 * @end points directly to the top most stack entry to avoid a -8
+	 * adjustment in the stack switch hotpath. Adjust it back before
+	 * calculating @begin.
+	 */
+	end++;
+	begin = end - (IRQ_STACK_SIZE / sizeof(long));
+
+	/*
+	 * Due to the switching logic RSP can never be == @end because the
+	 * final operation is 'popq %rsp' which means after that RSP points
+	 * to the original stack and not to @end.
+	 */
+	if (stack < begin || stack >= end)
+		return false;
+
+	info->type	= STACK_TYPE_REX;
+	info->begin	= begin;
+	info->end	= end;
+
+	/*
+	 * The next stack pointer is stored at the top of the irq stack
+	 * before switching to the irq stack. Actual stack entries are all
+	 * below that.
+	 */
+	info->next_sp = (unsigned long *)*(end - 1);
+
+	return true;
+}
+
 bool noinstr get_stack_info_noinstr(unsigned long *stack, struct task_struct *task,
 				    struct stack_info *info)
 {
@@ -185,6 +224,9 @@ bool noinstr get_stack_info_noinstr(unsigned long *stack, struct task_struct *ta
 	if (in_entry_stack(stack, info))
 		return true;
 
+	if (in_rex_stack(stack, info))
+		return true;
+
 	return false;
 }
 

arch/x86/kernel/smp.c

@@ -259,7 +259,7 @@ DEFINE_IDTENTRY_SYSVEC(sysvec_call_function)
 	apic_eoi();
 	trace_call_function_entry(CALL_FUNCTION_VECTOR);
 	inc_irq_stat(irq_call_count);
-	generic_smp_call_function_interrupt();
+	generic_smp_call_function_interrupt(regs);
 	trace_call_function_exit(CALL_FUNCTION_VECTOR);
 }
 
@@ -268,7 +268,7 @@ DEFINE_IDTENTRY_SYSVEC(sysvec_call_function_single)
 	apic_eoi();
 	trace_call_function_single_entry(CALL_FUNCTION_SINGLE_VECTOR);
 	inc_irq_stat(irq_call_count);
-	generic_smp_call_function_single_interrupt();
+	generic_smp_call_function_single_interrupt(regs);
 	trace_call_function_single_exit(CALL_FUNCTION_SINGLE_VECTOR);
 }
 

arch/x86/net/Makefile

@@ -8,3 +8,6 @@ ifeq ($(CONFIG_X86_32),y)
 else
         obj-$(CONFIG_BPF_JIT) += bpf_jit_comp.o bpf_timed_may_goto.o
 endif
+
+obj-y += rex.o
+obj-y += rex_$(BITS).o

arch/x86/net/rex.c

@@ -0,0 +1,105 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * X86-specific code for Rex support
+ */
+#define pr_fmt(fmt) "rex: " fmt
+
+#include <linux/bpf.h>
+#include <linux/compiler_types.h>
+#include <linux/module.h>
+#include <linux/objtool.h>
+#include <linux/percpu.h>
+#include <linux/printk.h>
+#include <linux/vmalloc.h>
+
+#include <asm/page.h>
+#include <asm/pgtable.h>
+#include <asm/rex.h>
+
+/* Align to page size, since the stack trace is broken anyway */
+struct rex_stack {
+	char stack[REX_STACK_SIZE];
+} __aligned(PAGE_SIZE);
+
+DEFINE_PER_CPU_PAGE_ALIGNED(struct rex_stack, rex_stack_backing_store)
+__visible;
+DEFINE_PER_CPU(void *, rex_stack_ptr);
+
+DECLARE_PER_CPU(const struct bpf_prog *, rex_curr_prog);
+
+/*
+ * Not supposed to be called by other kernel code, therefore keep prototype
+ * private
+ */
+void rex_landingpad(void) __noreturn;
+
+static int map_rex_stack(unsigned int cpu)
+{
+	char *stack = (char *)per_cpu_ptr(&rex_stack_backing_store, cpu);
+	struct page *pages[REX_STACK_SIZE / PAGE_SIZE];
+	void *va;
+	int i;
+
+	for (i = 0; i < REX_STACK_SIZE / PAGE_SIZE; i++) {
+		phys_addr_t pa = per_cpu_ptr_to_phys(stack + (i << PAGE_SHIFT));
+
+		pages[i] = pfn_to_page(pa >> PAGE_SHIFT);
+	}
+
+	va = vmap(pages, REX_STACK_SIZE / PAGE_SIZE, VM_MAP, PAGE_KERNEL);
+	if (!va)
+		return -ENOMEM;
+
+	/* Store actual TOS to avoid adjustment in the hotpath */
+	per_cpu(rex_stack_ptr, cpu) = va + REX_STACK_SIZE - 8;
+
+	pr_info("Initialize rex_stack on CPU %d at 0x%llx\n", cpu,
+		((u64)va) + REX_STACK_SIZE);
+
+	return 0;
+}
+
+int arch_init_rex_stack(void)
+{
+	int i, ret = 0;
+	for_each_online_cpu(i) {
+		ret = map_rex_stack(i);
+		if (ret < 0) {
+			pr_err("Failed to initialize rex stack on CPU %d\n", i);
+			break;
+		}
+	}
+	return ret;
+}
+
+/* Do not declare rex_landingpad_asm() in header file since it should only be
+ * called from rex_landingpad() 
+ */
+asmlinkage void __noreturn rex_landingpad_asm(void);
+
+void __noreturn rex_landingpad(void)
+{
+	struct task_struct *loader;
+	DEFINE_RATELIMIT_STATE(rex_rs, DEFAULT_RATELIMIT_INTERVAL,
+			       DEFAULT_RATELIMIT_BURST);
+
+	/* Report error */
+	if (__ratelimit(&rex_rs)) {
+		pr_err("%s\n", this_cpu_ptr(rex_log_buf));
+		dump_stack();
+	}
+
+	loader = find_task_by_pid_ns(
+		this_cpu_read_stable(rex_curr_prog)->saved_state->loader_pid,
+		&init_pid_ns);
+
+	/* Reuse the seccomp signal for now */
+	if (loader)
+		force_sig_fault_to_task(SIGSYS, SYS_SECCOMP, NULL, loader);
+
+	/* Reset the rex_termination_state set in rex panic handler */
+	this_cpu_write(rex_termination_state, 0);
+
+	/* Handle the rest fixups */
+	rex_landingpad_asm();
+}

arch/x86/net/rex_64.S

@@ -0,0 +1,84 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * X86 asm for Rex support
+ */
+#include <linux/export.h>
+#include <linux/linkage.h>
+#include <asm/errno.h>
+#include <asm/nospec-branch.h>
+#include <asm/percpu.h>
+
+	.code64
+	.section .text, "ax"
+
+/*
+ * Dispatcher func for Rex
+ *
+ * %rdi: ctx argument of the Rex program
+ * %rsi: pointer to the Rex program struct
+ * %rdx: fucntion pointer to the entry of the Rex program
+ */
+SYM_FUNC_START(rex_dispatcher_func)
+	/* save the callee-saved registers and the frame pointer*/
+	pushq %rbp
+	pushq %rbx
+	pushq %r12
+	pushq %r13
+	pushq %r14
+	pushq %r15
+
+	/* switch stack and save old rsp*/
+	movq PER_CPU_VAR(rex_stack_ptr), %rbp
+	movq %rsp, (%rbp)
+	movq %rbp, %rsp
+
+	/* record start time */
+	movq jiffies(%rip), %r11
+	movq %r11, PER_CPU_VAR(rex_prog_start_time)
+
+	/* let the timer know we are in */
+	movq %rsi, PER_CPU_VAR(rex_curr_prog)
+
+	/* invoke bpf func */
+	CALL_NOSPEC rdx
+
+/* 
+ * Exit path: rex_landingpad also redirects the control flow here
+ *
+ * %rax: program return value or default return value in case of a panic
+ * %rsp: top entry of the Rex stack
+ */
+SYM_INNER_LABEL(rex_exit, SYM_L_GLOBAL)
+	/* let the timer know we are out */
+	movq $0, PER_CPU_VAR(rex_curr_prog)
+
+	/* pop old stack pointer into rsp */
+	popq %rsp
+
+	/* restore the callee-saved registers and the frame pointer */
+	popq %r15
+	popq %r14
+	popq %r13
+	popq %r12
+	popq %rbx
+	popq %rbp
+
+	/* return */
+	RET
+SYM_FUNC_END(rex_dispatcher_func)
+EXPORT_SYMBOL(rex_dispatcher_func)
+
+/*
+ * Low-level fixups for Rust panics
+ */
+SYM_FUNC_START(rex_landingpad_asm)
+	/* set an return value of 0 */
+	movq $0, %rax
+
+	/* reset stack */
+	movq PER_CPU_VAR(rex_stack_ptr), %rsp
+
+	/* jump to exit path */
+	jmp rex_exit
+SYM_FUNC_END(rex_landingpad_asm)
+STACK_FRAME_NON_STANDARD(rex_landingpad_asm);

arch/x86/xen/smp.c

@@ -253,15 +253,15 @@ void xen_send_IPI_allbutself(int vector)
 
 static irqreturn_t xen_call_function_interrupt(int irq, void *dev_id)
 {
-	generic_smp_call_function_interrupt();
+	generic_smp_call_function_interrupt(dev_id);
 	inc_irq_stat(irq_call_count);
 
 	return IRQ_HANDLED;
 }
 
 static irqreturn_t xen_call_function_single_interrupt(int irq, void *dev_id)
 {
-	generic_smp_call_function_single_interrupt();
+	generic_smp_call_function_single_interrupt(dev_id);
 	inc_irq_stat(irq_call_count);
 
 	return IRQ_HANDLED;