#include <linux/module.h>
#include <linux/init.h>
#include <linux/kernel.h>   /* printk() */
#include <linux/slab.h>     /* kmalloc() */
#include <linux/fs.h>       /* everything... */
#include <linux/errno.h>    /* error codes */
#include <linux/types.h>    /* size_t */
#include <linux/cdev.h>     /* cdev */
#include <linux/fcntl.h>    /* O_ACCMODE */
#include <linux/uaccess.h>  /* copy_*_user */
#include <linux/version.h>  /* kernel version code */
#include <linux/moduleparam.h>

#if LINUX_VERSION_CODE < KERNEL_VERSION(5,0,0)
#define access_ok_wrapper(type,arg,cmd) access_ok(type, arg, cmd)
#else
#define access_ok_wrapper(type,arg,cmd) access_ok(arg, cmd)
#endif

/* ioctl definitions, use 'z' as magic number */
#define CHRDEV_IOC_MAGIC  'z'
#define CHRDEV_IOCRESET    _IO(CHRDEV_IOC_MAGIC, 0)
#define CHRDEV_IOCSET      _IOW(CHRDEV_IOC_MAGIC, 1, int)
#define CHRDEV_IOCGET      _IOR(CHRDEV_IOC_MAGIC, 2, int)
#define CHRDEV_IOC_MAXNR   3

/* device name and major number */
#define DEV_NAME         "chrdev"
#define DEV_SIZE         1024
#define CONFIG_AUTODEV   1 /* Auto create device node in driver or not */

//#define DEV_MAJOR  79
#ifndef DEV_MAJOR
#define DEV_MAJOR  0
#endif

int dev_major = DEV_MAJOR;
module_param(dev_major, int, S_IRUGO);

typedef struct chrdev_s
{
    struct cdev    cdev;
#ifdef CONFIG_AUTODEV
    struct class  *class;
    struct device *device;
#endif
    char          *data;   /* data buffer */
    uint32_t       size;   /* data buffer size */
    uint32_t       bytes;  /* data bytes in the buffer */
} chrdev_t;

static struct chrdev_s   dev;

static ssize_t chrdev_read (struct file *file, char __user *buf, size_t count, loff_t *f_pos)
{
    struct chrdev_s   *dev = file->private_data;
    ssize_t nbytes;
    ssize_t rv = 0;

    /* no data in buffer  */
    if( !dev->bytes )
        return 0;

    /* copy data to user space */
    nbytes = count>dev->bytes ? dev->bytes : count;
    if( copy_to_user(buf, dev->data, nbytes) )
    {
        rv = -EFAULT;
        goto out;
    }

    /* update return value and data bytes in buffer */
    rv = nbytes;
    dev->bytes -= nbytes;

out:
    return rv;
}

static ssize_t chrdev_write (struct file *file, const char __user *buf, size_t count, loff_t *f_pos)
{
    struct chrdev_s   *dev = file->private_data;
    ssize_t nbytes;
    ssize_t rv = 0;

    /* no space left */
    if( dev->bytes >= dev->size )
        return -ENOSPC;

    /* check copy data bytes */
    if( dev->size - dev->bytes < count)
        nbytes = dev->size - dev->bytes;
    else
        nbytes = count;

    /* copy data from user space  */
    if( copy_from_user(&dev->data[dev->bytes], buf, nbytes) )
    {
        rv = -EFAULT;
        goto out;
    }

    /* update return value and data bytes in buffer */
    rv = nbytes;
    dev->bytes += nbytes;

out:
    return rv;
}

/* The ioctl() implementation */
long chrdev_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
    struct chrdev_s *dev = file->private_data;
    int rv, data;
    size_t  bytes = sizeof(data);

    /*
     * extract the type and number bitfields, and don't decode
     * wrong cmds: return ENOTTY (inappropriate ioctl) before access_ok()
     */
    if (_IOC_TYPE(cmd) != CHRDEV_IOC_MAGIC) return -ENOTTY;
    if (_IOC_NR(cmd) > CHRDEV_IOC_MAXNR) return -ENOTTY;

    /*
     * the direction is a bitmask, and VERIFY_WRITE catches R/W transfers.
     * `Type' is user-oriented, while access_ok is kernel-oriented,
     * so the concept of "read" and "write" is reversed
     */
    if (_IOC_DIR(cmd) & _IOC_READ)
        rv = !access_ok_wrapper(VERIFY_WRITE, (void __user *)arg, _IOC_SIZE(cmd));
    else if (_IOC_DIR(cmd) & _IOC_WRITE)
        rv =  !access_ok_wrapper(VERIFY_READ, (void __user *)arg, _IOC_SIZE(cmd));
    if (rv) return -EFAULT;

    switch(cmd) {
        case CHRDEV_IOCRESET:
            dev->bytes = 0;
            memset(dev->data, 0, dev->size);
            rv = 0;
            break;

        /* Last 4 bytes in the buffer used to save ioctl() data*/
        case CHRDEV_IOCSET:
            rv = __get_user(data, (int __user *)arg);
            if( !rv )
                 memcpy(&dev->data[dev->size-bytes], &data, bytes);
            break;

        case CHRDEV_IOCGET:
            memcpy(&data, &dev->data[dev->size-bytes], bytes);
            rv = __put_user(data, (int __user *)arg);
            break;

        default:
            return -ENOTTY;
    }

    return rv;
}

static int chrdev_open (struct inode *inode, struct file *file)
{
    struct chrdev_s    *dev; /* device struct address */

    /* get the device struct address by container_of() */
    dev = container_of(inode->i_cdev, struct chrdev_s, cdev);

    /* save the device struct address for other methods */
    file->private_data = dev;

    return 0;
}

static int chrdev_close (struct inode *node, struct file *file)
{
    return 0;
}

static struct file_operations chrdev_fops = {
    .owner   = THIS_MODULE,
    .open    = chrdev_open,          /* open()  implementation */
    .read    = chrdev_read,          /* read()  implementation */
    .write   = chrdev_write,         /* write() implementation */
    .release = chrdev_close,         /* close() implementation */
    .unlocked_ioctl = chrdev_ioctl,  /* ioctl() implementation */
};

static int __init chrdev_init(void)
{
    dev_t      devno;
    int        rv;

    /* malloc and initial device read/write buffer */
    dev.data = kmalloc(DEV_SIZE, GFP_KERNEL);
    if( !dev.data )
    {
        printk(KERN_ERR " %s driver kmalloc() failed\n", DEV_NAME);
        return -ENOMEM;
    }
    dev.size = DEV_SIZE;
    dev.bytes = 0;
    memset(dev.data, 0, dev.size);

    /* dynamic alloc device node major number if not set */
    if(0 != dev_major)
    {
        devno = MKDEV(dev_major, 0);
        rv = register_chrdev_region(devno, 1, DEV_NAME);
    }
    else
    {
        rv = alloc_chrdev_region(&devno, 0, 1, DEV_NAME);
        dev_major = MAJOR(devno);
    }

    if(rv < 0)
    {
        printk(KERN_ERR "%s driver can't use major %d\n", DEV_NAME, dev_major);
        return -ENODEV;
    }

    /* setup and register cdev into kernel */
    cdev_init(&dev.cdev, &chrdev_fops);
    dev.cdev.owner = THIS_MODULE;
    rv = cdev_add(&dev.cdev, devno, 1);
    if( rv )
    {
        rv = -ENODEV;
        printk(KERN_ERR "%s driver regist failed, rv=%d\n", DEV_NAME, rv);
        goto failed1;
    }

#ifdef CONFIG_AUTODEV
    /* create device node in user space */
#if LINUX_VERSION_CODE >= KERNEL_VERSION(6, 5, 0)
    dev.class = class_create(DEV_NAME);
#else
    dev.class = class_create(THIS_MODULE, DEV_NAME);
#endif
    if (IS_ERR(dev.class)) {
        rv = PTR_ERR(dev.class);
        goto failed2;
    }

    dev.device = device_create(dev.class, NULL, MKDEV(dev_major, 0), NULL, DEV_NAME);
    if( !dev.device )
    {
        rv = -ENODEV;
        printk(KERN_ERR "%s driver create device failed\n", DEV_NAME);
        goto failed3;
    }
#endif

    printk(KERN_INFO "%s driver on major[%d] installed.\n", DEV_NAME, dev_major);
    return 0;

#ifdef CONFIG_AUTODEV
failed3:
    class_destroy(dev.class);

failed2:
    cdev_del(&dev.cdev);
#endif

failed1:
    unregister_chrdev_region(devno, 1);
    kfree(dev.data);

    printk(KERN_ERR "%s driver installed failed.\n", DEV_NAME);
    return rv;
}

static void __exit chrdev_exit(void)
{
#ifdef CONFIG_AUTODEV
    device_del(dev.device);
    class_destroy(dev.class);
#endif

    cdev_del(&dev.cdev);
    unregister_chrdev_region(MKDEV(dev_major,0), 1);

    kfree(dev.data);

    printk(KERN_INFO "%s driver removed!\n", DEV_NAME);
    return;
}

module_init(chrdev_init);
module_exit(chrdev_exit);

MODULE_LICENSE("Dual BSD/GPL");
MODULE_AUTHOR("GuoWenxue <guowenxue@gmail.com>");