CodeNote - Pass PC sampler

使用了有关项目的代码，请标记来源！

#include "wuwa_safe_signal.h"
#include "wuwa_common.h"
#include <linux/stddef.h>
#include <linux/signal_types.h>
#include <linux/ptrace.h>
#include <linux/kprobes.h>
#include <linux/vmalloc.h>

#include "wuwa_utils.h"

struct unsafe_region_area {
    u64 addr;
    u32 nums_page;
    uid_t uid;
    pid_t session_pid;
} __aligned(sizeof(u64));

static struct karray_list* unsafe_region_areas = NULL;
static DEFINE_RWLOCK(unsafe_region_areas_lock);

static bool is_safe_signal(struct ksignal *ksig, struct pt_regs *regs) {
    u64 lr = regs->regs[30];
    u64 pc = regs->pc;

    if (!unsafe_region_areas->data) {
        wuwa_err("unsafe region areas list is empty\n");
        return true;
    }

    read_lock(&unsafe_region_areas_lock);

    for (int i = 0; i < unsafe_region_areas->size; ++i) {
        struct unsafe_region_area *area = unsafe_region_areas->data[i];
        if (!area) {
            continue;
        }
        if (area->uid != current_uid().val) {
            continue;
        }
        u64 start = area->addr;
        u64 end = start + area->nums_page * PAGE_SIZE;
        if (pc >= start && pc < end) {
            // The PC is within a safe region
            wuwa_info("PC=0x%llx is in a unsafe region for pid %d\n", pc, current->pid);
            read_unlock(&unsafe_region_areas_lock);
            return false;
        }

        if (lr >= start && lr < end) {
            // The LR is within a safe region
            wuwa_info("LR=0x%llx is in a unsafe region for pid %d\n", lr, current->pid);
            read_unlock(&unsafe_region_areas_lock);
            return false;
        }
    }

    read_unlock(&unsafe_region_areas_lock);

    return true;
}

// static int setup_rt_frame(int usig, struct ksignal *ksig, sigset_t *set, struct pt_regs *regs)
static int handler_pre(struct kprobe *p, struct pt_regs *regs) {
    int usig = (int) regs->regs[0];
    struct ksignal *ksig = (struct ksignal *) regs->regs[1];
    sigset_t *set = (sigset_t *) regs->regs[2];
    if (usig < 0 || usig >= _NSIG) {
        wuwa_err("Invalid signal number: %d\n", usig);
        return 0;
    }
    if (!ksig || !set) {
        wuwa_err("Invalid ksignal or sigset_t pointer\n");
        return 0;
    }

    if (usig != SIGPROF) {
        return 0;
    }

    struct pt_regs *task_regs = (struct pt_regs *) regs->regs[3];
    if (is_safe_signal(ksig, task_regs)) {
        return 0;
    }

    regs->regs[0] = 1; // indicate that the signal is unsafe
    return 1;
}

static struct kprobe kp = {
    .symbol_name = "setup_rt_frame",
    .pre_handler = handler_pre,
    .post_handler = NULL,
};

int wuwa_safe_signal_init(void) {
    int ret;
    ret = register_kprobe(&kp);
    if (ret < 0) {
        wuwa_err("register_kprobe failed, returned %d\n", ret);
        return ret;
    }

    unsafe_region_areas = arraylist_create(ARRAYLIST_DEFAULT_CAPACITY);
    if (!unsafe_region_areas) {
        wuwa_err("failed to create unsafe region areas list\n");
        ret = -ENOMEM;
        goto out;
    }

    return 0;

    out:
    unregister_kprobe(&kp);
    return ret;
}

void wuwa_safe_signal_cleanup(void) {
    unregister_kprobe(&kp);
    write_lock(&unsafe_region_areas_lock);
    if (unsafe_region_areas) {
        if (!unsafe_region_areas->data) {
            write_unlock(&unsafe_region_areas_lock);
            wuwa_err("unsafe region areas list is empty\n");
            goto next;
        }

        for (int i = 0; i < unsafe_region_areas->size; ++i) {
            void* element = unsafe_region_areas->data[i];
            if (element)
                kvfree(element);
        }

        next:
        arraylist_destroy(unsafe_region_areas);
        unsafe_region_areas = NULL;
    }
    write_unlock(&unsafe_region_areas_lock);
}

int wuwa_add_unsafe_region(pid_t session, uid_t uid, uintptr_t start, size_t num_page) {
    write_lock(&unsafe_region_areas_lock);

    if (!unsafe_region_areas) {
        write_unlock(&unsafe_region_areas_lock);
        return -ENOMEM;
    }

    struct unsafe_region_area *area = kvzalloc(sizeof(*area), GFP_KERNEL);
    if (!area) {
        wuwa_err("failed to allocate memory for unsafe region area\n");
        write_unlock(&unsafe_region_areas_lock);
        return -ENOMEM;
    }
    area->addr = start;
    area->nums_page = num_page;
    area->uid = uid;
    area->session_pid = session;

    if (arraylist_add(unsafe_region_areas, area)) {
        write_unlock(&unsafe_region_areas_lock);
        wuwa_err("failed to add unsafe region area\n");
        return -ENOMEM;
    }

    wuwa_info("unsafe region area added for pid %d, start=0x%lx, num_page=%zu\n", session, start, num_page);
    write_unlock(&unsafe_region_areas_lock);
    return 0;
}

int wuwa_del_unsafe_region(pid_t pid) {
    write_lock(&unsafe_region_areas_lock);

    if (!unsafe_region_areas) {
        write_unlock(&unsafe_region_areas_lock);
        return -ENOMEM;
    }

    if (!unsafe_region_areas->data) {
        write_unlock(&unsafe_region_areas_lock);
        wuwa_err("unsafe region areas list is empty\n");
        return -ENOENT;
    }

    for (int i = 0; i < unsafe_region_areas->size; ++i) {
        struct unsafe_region_area *area = unsafe_region_areas->data[i];
        if (area && area->session_pid == pid) {
            void *removed = arraylist_remove(unsafe_region_areas, i);
            if (removed) {
                kvfree(removed);
            } else {
                wuwa_err("failed to remove unsafe region area for pid %d\n", pid);
            }
        }
    }

    write_unlock(&unsafe_region_areas_lock);
    return 0;
}