#include <linux/bio.h>
#include <linux/blk-mq.h>
#include <linux/blkdev.h>
#include <linux/errno.h>
#include <linux/fcntl.h>
#include <linux/fs.h>
#include <linux/genhd.h>
#include <linux/module.h>
#include <linux/string.h>
#include <linux/types.h>
#include <linux/vmalloc.h>

#define DISK_NAME "vramdisk"

#define MEMSIZE 0x19000 // Size of Ram disk in sectors

int major = 0; // Variable for Major Number

#define MDISK_SECTOR_SIZE 512
#define MBR_SIZE MDISK_SECTOR_SIZE
#define MBR_DISK_SIGNATURE_OFFSET 440
#define MBR_DISK_SIGNATURE_SIZE 4
#define PARTITION_TABLE_OFFSET 446
#define PARTITION_ENTRY_SIZE 16
#define PARTITION_TABLE_SIZE 64
#define MBR_SIGNATURE_OFFSET 510
#define MBR_SIGNATURE_SIZE 2
#define MBR_SIGNATURE 0xAA55
#define BR_SIZE MDISK_SECTOR_SIZE
#define BR_SIGNATURE_OFFSET 510
#define BR_SIGNATURE_SIZE 2
#define BR_SIGNATURE 0xAA55

typedef struct {
unsigned char boot_type; // 0x00 - Inactive; 0x80 - Active (Bootable)
unsigned char start_head;
unsigned char start_sec : 6;
unsigned char start_cyl_hi : 2;
unsigned char start_cyl;
unsigned char part_type;
unsigned char end_head;
unsigned char end_sec : 6;
unsigned char end_cyl_hi : 2;
unsigned char end_cyl;
unsigned int abs_start_sec;
unsigned int sec_in_part;
} PartEntry;

typedef PartEntry PartTable[4];

static struct lock_class_key sr_bio_compl_lkclass;

#define SEC_PER_HEAD 63
#define HEAD_PER_CYL 255
#define HEAD_SIZE (SEC_PER_HEAD * MDISK_SECTOR_SIZE)
#define CYL_SIZE (SEC_PER_HEAD * HEAD_PER_CYL * MDISK_SECTOR_SIZE)

#define sec4size(s) ((((s) % CYL_SIZE) % HEAD_SIZE) / MDISK_SECTOR_SIZE)
#define head4size(s) (((s) % CYL_SIZE) / HEAD_SIZE)
#define cyl4size(s) ((s) / CYL_SIZE)

static PartTable def_part_table =
{
{
boot_type: 0x00,
start_sec: 0x2,
start_head: 0x0,
start_cyl: 0x0,
part_type: 0x83,
end_head: 0x3,
end_sec: 0x20,
end_cyl: 0x9F,
abs_start_sec: 0x1,
sec_in_part: 0x4FFF // 10Mbyte
},
{
boot_type: 0x00,
start_head: 0x4,
start_sec: 0x1,
start_cyl: 0x0,
part_type: 0x05, // extended partition type
end_sec: 0x20,
end_head: 0xB,
end_cyl: 0x9F,
abs_start_sec: 0x5000,
sec_in_part: 0xA000
}
};

static unsigned int def_log_part_br_abs_start_sector[] = {0x5000, 0xA000};
static const PartTable def_log_part_table[] =
{
{
{
boot_type: 0x00,
start_head: 0x4,
start_sec: 0x2,
start_cyl: 0x0,
part_type: 0x83,
end_head: 0x7,
end_sec: 0x20,
end_cyl: 0x9F,
abs_start_sec: 0x1,
sec_in_part: 0x4FFF
},
{
boot_type: 0x00,
start_head: 0x8,
start_sec: 0x01,
start_cyl: 0x00,
part_type: 0x05,
end_head: 0xB,
end_sec: 0x20,
end_cyl: 0x9F,
abs_start_sec: 0x5000,
sec_in_part: 0x5000
}
},
{
{
boot_type: 0x00,
start_head: 0x8,
start_sec: 0x02,
start_cyl: 0x00,
part_type: 0x83,
end_head: 0xB,
end_sec: 0x20,
end_cyl: 0x9F,
abs_start_sec: 0x1,
sec_in_part: 0x4FFF
}
}
};

static void copy_mbr(u8 *disk) {
memset(disk, 0x0, MBR_SIZE);
*(unsigned long *)(disk + MBR_DISK_SIGNATURE_OFFSET) = 0x36E5756D;
memcpy(disk + PARTITION_TABLE_OFFSET, &def_part_table, PARTITION_TABLE_SIZE);
*(unsigned short *)(disk + MBR_SIGNATURE_OFFSET) = MBR_SIGNATURE;
}

static void copy_br(u8 *disk, int abs_start_sector,
const PartTable *part_table) {
disk += (abs_start_sector * MDISK_SECTOR_SIZE);
memset(disk, 0x0, BR_SIZE);
memcpy(disk + PARTITION_TABLE_OFFSET, part_table, PARTITION_TABLE_SIZE);
*(unsigned short *)(disk + BR_SIGNATURE_OFFSET) = BR_SIGNATURE;
}

void copy_mbr_n_br(u8 *disk) {
int i;
copy_mbr(disk);
for (i = 0; i < ARRAY_SIZE(def_log_part_table); ++i)
copy_br(disk, def_log_part_br_abs_start_sector[i], &def_log_part_table[i]);
}

/* Structure associated with Block device*/
struct vram_disk_device {
size_t size;
u8 *data;
atomic_t nr_users;
// spinlock_t lock;
struct blk_mq_tag_set tag_set;
struct request_queue *queue;
struct gendisk *gd;
} device;

// struct vram_disk_device *x;

static int vramdisk_open(struct block_device *x, fmode_t mode) {
struct vram_disk_device *dev = x->bd_disk->private_data;
if (dev == NULL) {
printk(KERN_WARNING DISK_NAME "driver: invalid disk private_data\n");
return -ENXIO;
}

atomic_inc(&dev->nr_users);
printk(KERN_INFO DISK_NAME "driver: open \n");

return 0;
}

static void vramdisk_release(struct gendisk *disk, fmode_t mode) {
struct vram_disk_device *dev = disk->private_data;
if (dev) {
atomic_dec(&dev->nr_users);
printk(KERN_INFO DISK_NAME "driver: closed \n");
} else
printk(KERN_WARNING DISK_NAME "driver: invalid disk private_data\n");
}

static struct block_device_operations fops = {
.owner = THIS_MODULE,
.open = vramdisk_open,
.release = vramdisk_release,
};

int vramdisk_init(void) {
device.size = MEMSIZE;
(device.data) = vmalloc(MEMSIZE * MDISK_SECTOR_SIZE);
if (device.data == NULL) {
printk(KERN_WARNING DISK_NAME "driver: vmalloc failure.\n");
return -ENOMEM;
}

/* Load partition table */
copy_mbr_n_br(device.data);

return 0;
}

static int rb_transfer(struct request *req, unsigned int *nr_bytes) {
int dir = rq_data_dir(req);
sector_t start_sector = blk_rq_pos(req);
unsigned int sector_cnt =
blk_rq_sectors(req); /* no of sector on which opn to be done*/

struct bio_vec bv;
#define BV_PAGE(bv) ((bv).bv_page)
#define BV_OFFSET(bv) ((bv).bv_offset)
#define BV_LEN(bv) ((bv).bv_len)

struct req_iterator iter;
sector_t sector_offset = 0;
unsigned int sectors;
u8 *buffer;

rq_for_each_segment(bv, req, iter) {
buffer = page_address(BV_PAGE(bv)) + BV_OFFSET(bv);
if (BV_LEN(bv) % (MDISK_SECTOR_SIZE) != 0) {
printk(KERN_ERR DISK_NAME
": bio size is not a multiple of sector size\n");
return -EIO;
}

sectors = BV_LEN(bv) / MDISK_SECTOR_SIZE;
printk(KERN_DEBUG DISK_NAME ": Start Sector: %llu, Sector Offset: %llu; "
"Buffer: %p; Length: %u sectors\n",
(unsigned long long)(start_sector),
(unsigned long long)(sector_offset), buffer, sectors);

if (dir == WRITE) /* Write to the device */
memcpy((device.data) +
((start_sector + sector_offset) * MDISK_SECTOR_SIZE),
buffer, sectors * MDISK_SECTOR_SIZE);
else /* Read from the device */
memcpy(buffer,
(device.data) +
((start_sector + sector_offset) * MDISK_SECTOR_SIZE),
sectors * MDISK_SECTOR_SIZE);

sector_offset += sectors;
*nr_bytes += sectors * MDISK_SECTOR_SIZE;
}

if (sector_offset != sector_cnt) {
printk(KERN_ERR DISK_NAME ": bio info doesn't match with the request info");
return -EIO;
}

return 0;
#undef BV_PAGE
#undef BV_OFFSET
#undef BV_LEN
}

/* request handling function */
/**
* Changing according to blk multi-queue feature
*/
static blk_status_t handle_request(struct blk_mq_hw_ctx *hardware_context,
const struct blk_mq_queue_data *bd) {
blk_status_t status = BLK_STS_OK;
unsigned int nr_bytes = 0;
struct request *req = bd->rq;
switch (rb_transfer(req, &nr_bytes)) {
case -EIO: {
status = BLK_STS_IOERR;
} break;
default: {
status = BLK_STS_OK;
} break;
};

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

return status;
}

static struct blk_mq_ops vram_queue_ops = {
.queue_rq = handle_request,
};

int device_setup(void) {

/* Register block device */
major = register_blkdev(major, DISK_NAME); // major no. allocation

if (major <= 0) {
printk(KERN_WARNING DISK_NAME ": unable to get major number\n");
return -EBUSY;
}

printk(KERN_ALERT "Major Number is : %d", major);

if (vramdisk_init()) {
return -ENOMEM;
}

device.tag_set.ops = &vram_queue_ops;
device.tag_set.nr_hw_queues = 1;
device.tag_set.queue_depth = 128;
device.tag_set.numa_node = NUMA_NO_NODE;
device.tag_set.cmd_size = sizeof(struct vram_disk_device);
device.tag_set.flags = BLK_MQ_F_SHOULD_MERGE;
device.tag_set.driver_data = &device;
if (blk_mq_alloc_tag_set(&(device.tag_set))) {
return -ENOMEM;
}

device.queue = blk_mq_init_queue(&(device.tag_set));
if (IS_ERR(device.queue)) {
blk_mq_free_tag_set(&(device.tag_set));
return -ENOMEM;
}
device.queue->queuedata = &device;

/* Initialize gendisk */
device.gd = __alloc_disk_node(device.queue, 7, &sr_bio_compl_lkclass);

if (!device.gd) {
printk(KERN_INFO DISK_NAME ": alloc_disk failure\n");
return -EBUSY;
}

device.gd->minors = 7;
(device.gd)->major = major; // major no to gendisk
device.gd->first_minor = 0; // first minor of gendisk

device.gd->fops = &fops;
device.gd->private_data = &device;
device.gd->queue = device.queue;
printk(KERN_INFO "THIS IS DEVICE SIZE %zu", device.size);
//
snprintf(((device.gd)->disk_name), DISK_NAME_LEN, DISK_NAME);

set_capacity(device.gd, device.size);

add_disk(device.gd);
return 0;
}

static int __init vramdisk_drive_init(void) {
int ret = 0;
ret = device_setup();
return ret;
}

void vramdisk_cleanup(void) { vfree(device.data); }

void __exit vramdisk_drive_exit(void) {
del_gendisk(device.gd);
put_disk(device.gd);

blk_cleanup_queue(device.queue);
blk_mq_free_tag_set(&(device.tag_set));

// cleanup device buffer in ram
vramdisk_cleanup();

unregister_blkdev(major, DISK_NAME);
}

module_init(vramdisk_drive_init);
module_exit(vramdisk_drive_exit);
MODULE_LICENSE("Dual MIT/GPL");
MODULE_AUTHOR("Bitway");
MODULE_DESCRIPTION("BLOCK DRIVER");