Simple ramdisk block device kernel module for linux kernel 6

ramdisk.c

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kthread.h>
#include <linux/time.h>
#include <linux/timer.h>
#include <linux/fs.h>
#include <asm/segment.h>
#include <asm/uaccess.h>
#include <linux/buffer_head.h>
#include <linux/vmalloc.h>
#include <linux/slab.h>
#include <linux/blkdev.h>
#include <linux/hdreg.h>
#include <linux/blk-mq.h>
#include <linux/delay.h>

#define DEVICENAME "ramdisk"
#define BLOCKDEVICE_NAME "ramdisk0"

#define RAMDISK_MEM_SIZE 1073741824 // 1GB
#define DISK_NUM_SECTORS (RAMDISK_MEM_SIZE / 512)

static int init_done = 0;
static int disk_initialized = 0;
static unsigned char *ramdisk_mem = NULL;

static int major_dev = -1;

struct f_device_s {
    sector_t dev_num_sectors;
    struct blk_mq_tag_set tag_set;
    struct gendisk *disk;
};

static blk_status_t queue_rq(struct blk_mq_hw_ctx *hctx, const struct blk_mq_queue_data *bd);

static struct blk_mq_ops _mq_ops = {
    .queue_rq = queue_rq
};

struct f_device_s *f_device = NULL;

static int fops_open(struct block_device *, fmode_t) {
    return 0;
}
static void fops_release(struct gendisk *, fmode_t) {
}
static int fops_ioctl(struct block_device *, fmode_t, unsigned, unsigned long) {
    return 0;
}

int fops_compat_ioctl(struct block_device *, fmode_t, unsigned, unsigned long) {\
    return 0;
}

static const struct block_device_operations _fops = {
    .owner = THIS_MODULE,
    .open = fops_open,
    .release = fops_release,
    .ioctl = fops_ioctl,
    #ifdef CONFIG_COMPAT
        .compat_ioctl = fops_compat_ioctl,
    #endif
};

static int do_simple_request(struct request * rq, unsigned int * nr_bytes)
{
    int ret = 0;
    struct bio_vec bvec;
    struct req_iterator iter;
    struct f_device_s *dev = rq->q->queuedata;
    loff_t pos = blk_rq_pos(rq) /* current sector */ << SECTOR_SHIFT;
    loff_t dev_size = (loff_t)(dev->dev_num_sectors << SECTOR_SHIFT);

    rq_for_each_segment(bvec, rq, iter)
    {
        unsigned long b_len = bvec.bv_len;

        void *b_buf = page_address(bvec.bv_page) + bvec.bv_offset;

        if ((pos + b_len) > dev_size) {
            b_len = (unsigned long) (dev_size - pos);
        }

        if (rq_data_dir (rq)) // WRITE
        {
            memcpy(ramdisk_mem + pos, b_buf, b_len);
        }
        else // READ
        {
            memcpy(b_buf, ramdisk_mem + pos, b_len);
        }

        pos += b_len;
        *nr_bytes += b_len;
    }

    return ret;
}

blk_status_t queue_rq(struct blk_mq_hw_ctx *hctx, const struct blk_mq_queue_data *bd) {
    blk_status_t status = BLK_STS_OK;
    struct request *rq = bd->rq;

    blk_mq_start_request (rq);

    // we can't use that thread
    {
        int nr_bytes = 0;

        if (do_simple_request(rq, &nr_bytes) != 0)
            status = BLK_STS_IOERR;

        if (blk_update_request (rq, status, nr_bytes)) // GPL-only symbol
            BUG ();
        __blk_mq_end_request (rq, status);
    }

    return BLK_STS_OK;
}

static void setup_f_device(void) {
    memset(f_device, 0, sizeof(struct f_device_s));

    f_device->dev_num_sectors = DISK_NUM_SECTORS;

    f_device->tag_set.ops = &_mq_ops;
    f_device->tag_set.nr_hw_queues = 1;
    f_device->tag_set.queue_depth = 128;
    f_device->tag_set.numa_node = NUMA_NO_NODE;
    f_device->tag_set.cmd_size = 8;
    f_device->tag_set.flags = BLK_MQ_F_SHOULD_MERGE/* | BLK_MQ_F_SG_MERGE*/; //TODO:
    f_device->tag_set.driver_data = f_device;
    
    if (blk_mq_alloc_tag_set(&f_device->tag_set) != 0) {
        printk(KERN_WARNING "Failed to alloc tag set\n");
        return;
    }

    // disk
    {
        static struct lock_class_key __key;
        struct gendisk *disk = __blk_mq_alloc_disk(&f_device->tag_set, f_device,
		                            &__key);
        disk->flags |= GENHD_FL_NO_PART;
        disk->flags |= GENHD_FL_REMOVABLE;
        disk->major = major_dev;
        disk->minors = 1;
        disk->first_minor = 0;
        disk->fops = &_fops;
        disk->private_data = f_device;
        sprintf(disk->disk_name, BLOCKDEVICE_NAME);
        set_capacity(disk, f_device->dev_num_sectors);

        f_device->disk = disk;

        if (add_disk(disk) != 0) {
            printk(KERN_WARNING "failed to initialize disk\n");
        } else {
            disk_initialized = 1;
        }
    }
}

int init_module(void) {
    printk(KERN_INFO "Ramdisk kernel module\n");

    ramdisk_mem = vmalloc(RAMDISK_MEM_SIZE);

    if (!ramdisk_mem) {
        printk(KERN_INFO "failed to alloc key\n");
        return 0;
    }

    major_dev = register_blkdev(0, DEVICENAME);

    if (major_dev <= 0) {
        printk(KERN_ALERT KERN_INFO "Could not get a major number!\n");
        return -EBUSY;
    }

    f_device = kmalloc(sizeof(struct f_device_s), GFP_KERNEL);

    if(!f_device)
    {
        unregister_blkdev(major_dev, DEVICENAME);
        return -ENOMEM;
    }

    setup_f_device();

    return 0;
}

void cleanup_module(void) {
    if (disk_initialized) {
        del_gendisk(f_device->disk);
        blk_mq_free_tag_set(&f_device->tag_set);
    }

    if (major_dev > 0) {
        unregister_blkdev(major_dev, DEVICENAME);
    }

    if (ramdisk_mem) {
        vfree(ramdisk_mem);
    }

    kfree(f_device);
}

MODULE_LICENSE("GPL");

Based on https://prog.world/linux-kernel-5-0-we-write-simple-block-device-under-blk-mq/