#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/bio.h>
#include <linux/bvec.h>
#include <linux/blkdev.h>
#include <linux/blk-mq.h>

static int submit_bio_pre(struct kprobe *p, struct pt_regs *regs) {
    struct bio *bio = (struct bio *)regs->di;
    struct bio_vec bvec;
    struct bvec_iter iter;

    printk(KERN_INFO "submit_bio: bi_vec: bi_vcnt is %d, operation is %s\n", bio->bi_vcnt, blk_op_str(bio_op(bio)));
    bio_for_each_bvec(bvec, bio, iter) {
        printk(KERN_INFO "    idx=%d, sector=%lld, size=%d, done=%d, bv_page=0x%lx, bv_len=%u, bv_offset=%u\n",
               iter.bi_idx, iter.bi_sector, iter.bi_size, iter.bi_bvec_done, (unsigned long)bvec.bv_page, bvec.bv_len, bvec.bv_offset);
    }

    return 0;
}

static struct kprobe submit_bio_kp = {
    .symbol_name = "submit_bio",
    .pre_handler = submit_bio_pre,
};

static int blk_mq_start_request_pre(struct kprobe *p, struct pt_regs *regs) {
    struct request *rq = (struct request *)regs->di;

    printk(KERN_INFO "tag is %d, internal_tag is %d, sector is %lld, bytes is %d, nr_phys_segments is %d\n", rq->tag, rq->internal_tag, blk_rq_pos(rq), blk_rq_bytes(rq), rq->nr_phys_segments);

    return 0;
}

static struct kprobe blk_mq_start_request_kp = {
    .symbol_name = "blk_mq_start_request",
    .pre_handler = blk_mq_start_request_pre,
};

static struct kprobe *kprobes[] = {
    &submit_bio_kp,
    &blk_mq_start_request_kp,
};

static int __init kprobe_init(void) {
    int ret;

    ret = register_kprobes(kprobes, ARRAY_SIZE(kprobes));
    if (ret < 0) {
        printk(KERN_INFO "kprobe registration failed\n");
        return ret;
    }
    printk(KERN_INFO "kprobemodule loaded\n");

    return 0;
}

static void __exit kprobe_exit(void) {
    unregister_kprobes(kprobes, ARRAY_SIZE(kprobes));
    printk(KERN_INFO "kprobemodule unloaded\n");
}

module_init(kprobe_init);
module_exit(kprobe_exit);
MODULE_LICENSE("GPL");