Linux kernel module for Zen workaround for rr

Makefile

obj-m = zen_workaround.o

all:
	make -C /lib/modules/$(shell uname -r)/build M=$(PWD) modules

clean:
	make -C /lib/modules/$(shell uname -r)/build M=$(PWD) clean

zen_workaround.c

#include <linux/module.h>
#include <linux/kallsyms.h>

#define MODULE_NAME "zen_workaround"

typedef int set_memory_fn_t(unsigned long, int);

unsigned long need_symbol(const char *name)
{
	unsigned long sym = kallsyms_lookup_name(name);
	if (sym) {
		pr_info("%s is 0x%lx\n", name, sym);
	} else {
		pr_err("could not find kernel symbol %s\n", name);
	}
	return sym;
}

static set_memory_fn_t *set_memory_ro, *set_memory_rw;
static _Atomic u64 *x86_amd_ls_cfg_base_ptr;
static _Atomic int *kernel_set_to_readonly_ptr;

u64 set_speclockmap_disable(u64 msr) {
	return msr | BIT_64(54);
}

u64 unset_speclockmap_disable(u64 msr) {
	return msr & ~BIT_64(54);
}

typedef u64 (*edit_msr_func_t)(u64);

static void edit_ls_cfg_on_cpu(void *info)
{
	int cpu = get_cpu();
	u64 value = 0;

	if (!rdmsrl_safe(MSR_AMD64_LS_CFG, &value)) {
		edit_msr_func_t edit_msr = (edit_msr_func_t) info;
		u64 new_value = edit_msr(value);
		if (!wrmsrl_safe(MSR_AMD64_LS_CFG, new_value)) {
			pr_info("MSR_AMD64_LS_CFG for cpu %d was 0x%llx, setting to 0x%llx\n",
			        cpu, value, new_value);
		} else {
			pr_err("MSR_AMD64_LS_CFG for cpu %d was 0x%llx, setting to 0x%llx failed\n",
			       cpu, value, new_value);
		}
	}
}

static int do_zen_workaround(edit_msr_func_t edit_msr)
{
	if (boot_cpu_has(X86_FEATURE_LS_CFG_SSBD)) {
		int ret;
		u64 old_value = *x86_amd_ls_cfg_base_ptr;
		u64 new_value = edit_msr(old_value);

		*kernel_set_to_readonly_ptr = 0;
		ret = set_memory_rw((unsigned long)x86_amd_ls_cfg_base_ptr, 1);
		*kernel_set_to_readonly_ptr = 1;

		if (ret) {
			pr_err("set_memory_rw failed with %d\n", ret);
			return -EPERM;
		}

		*x86_amd_ls_cfg_base_ptr = new_value;
		pr_info("x86_amd_ls_cfg_base was 0x%llx, setting to 0x%llx\n", old_value, new_value);

		*kernel_set_to_readonly_ptr = 0;
		ret = set_memory_ro((unsigned long)x86_amd_ls_cfg_base_ptr, 1);
		*kernel_set_to_readonly_ptr = 1;

		if (ret) {
			pr_err("set_memory_ro failed with %d\n", ret);
			// Not returning early because we still presumably set x86_amd_ls_cfg_base.
		}
	} else {
		pr_info("x86_amd_ls_cfg_base is unused.");
	}

	smp_call_function(edit_ls_cfg_on_cpu, edit_msr, 1);
	edit_ls_cfg_on_cpu(edit_msr);
	return 0;
}

static int __init zen_workaround_init(void)
{
	if (!static_cpu_has(X86_FEATURE_ZEN)) {
		pr_err("Cannot use the Zen workaround on a non-Zen CPU\n");
		return -EINVAL;
	}

	set_memory_ro = (set_memory_fn_t*) need_symbol("set_memory_ro");
	set_memory_rw = (set_memory_fn_t*) need_symbol("set_memory_rw");
	x86_amd_ls_cfg_base_ptr = (_Atomic u64*) need_symbol("x86_amd_ls_cfg_base");
	kernel_set_to_readonly_ptr = (_Atomic int*) need_symbol("kernel_set_to_readonly");

	if (!set_memory_ro || !set_memory_rw || !x86_amd_ls_cfg_base_ptr ||
		      	!kernel_set_to_readonly_ptr) {
		return -ENOENT;
	}

	return do_zen_workaround(set_speclockmap_disable);
}
module_init(zen_workaround_init);

static void __exit zen_workaround_exit(void)
{
	do_zen_workaround(unset_speclockmap_disable);
}
module_exit(zen_workaround_exit)

MODULE_LICENSE("GPL");

Inspired by https://gist.github.com/eddyb/b888bb87988ca97ead9abcf96aa49e15

outdated comment, click to open

(only applies to https://gist.github.com/glandium/01d54cefdb70561b5f6675e08f2990f2/6147e24ad62ba3b3023eabb6aedfd0bd592839da)

Note that you could make this work using System.map for all symbols as long as you can import at least one symbol directly, I only used kallsyms_lookup_name for convenience.

EDIT: more details in rr-debugger/rr#2034 (comment)

Latest version avoids requiring internal kernel symbols altogether.

On NixOS, as per https://nixos.wiki/wiki/Linux_kernel#Developing_kernel_modules, one should be able to just:

make -C $(nix-build -E '(import <nixpkgs> {}).linux.dev' --no-out-link)/lib/modules/*/build M=$(pwd) modules

The Makefile in the gist is fine to use because the obj-m = zen_workaround.o line will be used and the rest ignored.

In case you have SecureBoot enabled (tested on Ubuntu):

sign: $(obj-m)
        /lib/modules/$(shell uname -r)/build/scripts/sign-file sha512 /var/lib/shim-signed/mok/MOK.priv /var/lib/shim-signed/mok/MOK.der zen_workaround.ko

I have packaged this into a NixOS module.

Kernel 6.5.8 on Ubuntu 23.10:

[ 9058.474358] BUG: scheduling while atomic: swapper/15/0/0x00000002
[ 9058.474360] Modules linked in:
[ 9058.474360] BUG: scheduling while atomic: swapper/5/0/0x00000002
[ 9058.474361]  zen_workaround(O+)
[ 9058.474363] Modules linked in: zen_workaround(O+)

Fix:

--- zen_workaround.c.orig       2023-11-06 16:37:23.000582213 +0100
+++ zen_workaround.c    2023-11-06 16:37:24.628603272 +0100
@@ -32,6 +32,8 @@
                               cpu, value, new_value);
                }
        }
+
+       put_cpu();
 }

 static void do_zen_workaround(edit_msr_func_t edit_msr)

Fix:

--- zen_workaround.c.orig       2023-11-06 16:37:23.000582213 +0100
+++ zen_workaround.c    2023-11-06 16:37:24.628603272 +0100
@@ -32,6 +32,8 @@
                               cpu, value, new_value);
                }
        }
+
+       put_cpu();
 }

 static void do_zen_workaround(edit_msr_func_t edit_msr)

Applied. Thanks.

Could this module be moved to a proper repository? this would make it easier to track updates.

Maybe it should just be in the rr repo. Open an issue?

rr-debugger/rr#3690