Vmx/hypervisor/vmm exploitation fun

#include <linux/cpumask.h>
#include <linux/smp.h>
#include <linux/slab.h>
#include <linux/version.h>
#if LINUX_VERSION_CODE <= KERNEL_VERSION(4, 8, 0)
#include <linux/cpu.h>
#include <linux/notifier.h>
#else
#include <linux/cpuhotplug.h>
#endif
#include <asm/special_insns.h>
#include <asm/cpufeature.h>

#include "cpu.h"
#include "crx.h"
#include "msr.h"
#include "mem.h"
#include "asm.h"
#include "enc.h"

#define VCPU_DBG(fmt, ...) \
    pr_debug("vcpu [cpu %02d] - " fmt, cur_logical_cpu(), ##__VA_ARGS__)

#define OLD_CPUHOTPLUG (LINUX_VERSION_CODE <= KERNEL_VERSION(4, 8, 0))

// Callback for CPU online
static int cpu_on_cb(unsigned int cpu)
{
    // Placeholder for vCPU virtualization
    return 0;
}

// Callback for CPU offline
static int cpu_off_cb(unsigned int cpu)
{
    // Placeholder for restoring vCPU and cleaning up data
    return 0;
}

#if OLD_CPUHOTPLUG
static int cpu_hotplug_cb(struct notifier_block *nblock,
                          unsigned long action,
                          void *hcpu)
{
    unsigned int cpu = (unsigned int)(unsigned long)hcpu;

    switch (action) {
    case CPU_ONLINE:
        return cpu_on_cb(cpu);
    case CPU_OFFLINE:
        return cpu_off_cb(cpu);
    default:
        return NOTIFY_DONE;
    }
}

static struct notifier_block cpu_hotplug_notifier = {
    .notifier_call = cpu_hotplug_cb,
};
#endif

int cpu_hotplug_register(void)
{
    int ret = 1;

#if OLD_CPUHOTPLUG
    register_cpu_notifier(&cpu_hotplug_notifier);
#else
    ret = cpuhp_setup_state_nocalls(CPUHP_AP_ONLINE_DYN,
                                    "vcpu:online",
                                    cpu_on_cb,
                                    cpu_off_cb);
    if (ret < 0) {
        pr_err("Failed to register CPU hotplug callback: %d\n", ret);
        return 0;
    }
#endif

    VCPU_DBG("Registered CPU hotplug callback");
    return ret;
}

void cpu_hotplug_unregister(void)
{
#if OLD_CPUHOTPLUG
    unregister_cpu_notifier(&cpu_hotplug_notifier);
#else
    cpuhp_remove_state_nocalls(CPUHP_AP_ONLINE_DYN);
#endif

    VCPU_DBG("Unregistered CPU hotplug callback");
}

int max_logical_cpu(void)
{
    return num_online_cpus();
}

int cur_logical_cpu(void)
{
    return smp_processor_id();
}

static inline bool vcpu_supports_vmx(void)
{
    return this_cpu_has(X86_FEATURE_VMX);
}

static int vcpu_set_feature_control(void)
{
    ia32_feature_control_t feat_ctl = {0};

    feat_ctl.value = __rdmsr(IA32_FEATURE_CONTROL_MSR);

    if (feat_ctl.fields.locked) {
        if (!feat_ctl.fields.vmx_outside_smx) {
            VCPU_DBG("VMX operation disabled by BIOS/UEFI firmware");
            return -EINVAL;
        }
        VCPU_DBG("VMX operation already enabled");
        return 0;
    }

    feat_ctl.fields.locked = 1;
    feat_ctl.fields.vmx_outside_smx = 1;
    __wrmsr(IA32_FEATURE_CONTROL_MSR, (feat_ctl.value >> 32), feat_ctl.value);

    feat_ctl.value = __rdmsr(IA32_FEATURE_CONTROL_MSR);

    if (!feat_ctl.fields.locked || !feat_ctl.fields.vmx_outside_smx) {
        VCPU_DBG("Failed to enable VMX operation");
        return -EIO;
    }

    return 0;
}

static void vcpu_set_cr_fixed(void)
{
    cr0_t cr0 = {0};
    cr4_t cr4 = {0};
    ia32_gen_fixed_t fixed = {0};

    cr4.value = __read_cr4();
    cr4.flags.VMXE = 1;
    __write_cr4(cr4.value);

    cr0.value = __read_cr0();
    fixed.value = __rdmsr(IA32_VMX_CR0_FIXED0_MSR);
    cr0.value |= fixed.split.low;
    fixed.value = __rdmsr(IA32_VMX_CR0_FIXED1_MSR);
    cr0.value &= fixed.split.low;
    __write_cr0(cr0.value);

    cr4.value = __read_cr4();
    fixed.value = __rdmsr(IA32_VMX_CR4_FIXED0_MSR);
    cr4.value |= fixed.split.low;
    fixed.value = __rdmsr(IA32_VMX_CR4_FIXED1_MSR);
    cr4.value &= fixed.split.low;
    __write_cr4(cr4.value);
}

vcpu_ctx_t *vcpu_alloc(void)
{
    vcpu_ctx_t *vcpu_ctx = kmalloc(sizeof(*vcpu_ctx), GFP_KERNEL);

    if (!vcpu_ctx) {
        pr_err("Failed to allocate vCPU context\n");
        return NULL;
    }

    vcpu_ctx->vmxon_region = mem_alloc_pages(0);
    if (!vcpu_ctx->vmxon_region) {
        pr_err("Failed to allocate VMXON region\n");
        kfree(vcpu_ctx);
        return NULL;
    }
    vcpu_ctx->vmxon_physical = mem_virt_to_phys(vcpu_ctx->vmxon_region);

    vcpu_ctx->vmcs_region = mem_alloc_pages(0);
    if (!vcpu_ctx->vmcs_region) {
        pr_err("Failed to allocate VMCS region\n");
        mem_free_pages((unsigned long)vcpu_ctx->vmxon_region, 0);
        kfree(vcpu_ctx);
        return NULL;
    }
    vcpu_ctx->vmcs_physical = mem_virt_to_phys(vcpu_ctx->vmcs_region);

    ia32_vmx_basic_t basic = {0};
    basic.value = __rdmsr(IA32_VMX_BASIC_MSR);
    vcpu_ctx->vmxon_region->reserved.rev_ident = basic.fields.vmcs_rev_ident;
    vcpu_ctx->vmcs_region->reserved.rev_ident = basic.fields.vmcs_rev_ident;

    return vcpu_ctx;
}

void vcpu_free(vcpu_ctx_t *vcpu_ctx)
{
    if (!vcpu_ctx)
        return;

    if (vcpu_ctx->vmxon_region)
        mem_free_pages((unsigned long)vcpu_ctx->vmxon_region, 0);

    if (vcpu_ctx->vmcs_region)
        mem_free_pages((unsigned long)vcpu_ctx->vmcs_region, 0);

    kfree(vcpu_ctx);
}

void vcpu_init(void *info)
{
    vmm_ctx_t *vmm_ctx = info;
    int cpu = cur_logical_cpu();

    if (!vcpu_supports_vmx()) {
        VCPU_DBG("VMX not supported");
        return;
    }

    if (vcpu_set_feature_control()) {
        VCPU_DBG("Failed to set IA32_FEATURE_CONTROL MSR");
        return;
    }

    vcpu_set_cr_fixed();

    vcpu_ctx_t *vcpu_ctx = vmm_ctx->vcpu_ctxs[cpu];

    if (__vmx_on(vcpu_ctx->vmxon_physical)) {
        VCPU_DBG("VMXON failed");
        return;
    }

    VCPU_DBG("VMXON successful");
}

void vcpu_restore(void *info)
{
    __vmx_off();
    VCPU_DBG("VMXOFF executed");
}