// SPDX-License-Identifier: GPL-2.0 OR Linux-OpenIB
/*
 * Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved
 */

#include <linux/kref.h>
#include <linux/cdev.h>
#include <linux/mutex.h>
#include <linux/file.h>
#include <linux/fs.h>
#include <rdma/ib_ucaps.h>

#define RDMA_UCAP_FIRST RDMA_UCAP_MLX5_CTRL_LOCAL

static DEFINE_MUTEX(ucaps_mutex);
static struct ib_ucap *ucaps_list[RDMA_UCAP_MAX];
static bool ucaps_class_is_registered;
static dev_t ucaps_base_dev;

struct ib_ucap {
        struct cdev cdev;
        struct device dev;
        struct kref ref;
};

static const char *ucap_names[RDMA_UCAP_MAX] = {
        [RDMA_UCAP_MLX5_CTRL_LOCAL] = "mlx5_perm_ctrl_local",
        [RDMA_UCAP_MLX5_CTRL_OTHER_VHCA] = "mlx5_perm_ctrl_other_vhca"
};

static char *ucaps_devnode(const struct device *dev, umode_t *mode)
{
        if (mode)
                *mode = 0600;

        return kasprintf(GFP_KERNEL, "infiniband/%s", dev_name(dev));
}

static const struct class ucaps_class = {
        .name = "infiniband_ucaps",
        .devnode = ucaps_devnode,
};

static const struct file_operations ucaps_cdev_fops = {
        .owner = THIS_MODULE,
        .open = simple_open,
};

/**
 * ib_cleanup_ucaps - cleanup all API resources and class.
 *
 * This is called once, when removing the ib_uverbs module.
 */
void ib_cleanup_ucaps(void)
{
        mutex_lock(&ucaps_mutex);
        if (!ucaps_class_is_registered) {
                mutex_unlock(&ucaps_mutex);
                return;
        }

        for (int i = RDMA_UCAP_FIRST; i < RDMA_UCAP_MAX; i++)
                WARN_ON(ucaps_list[i]);

        class_unregister(&ucaps_class);
        ucaps_class_is_registered = false;
        unregister_chrdev_region(ucaps_base_dev, RDMA_UCAP_MAX);
        mutex_unlock(&ucaps_mutex);
}

static int get_ucap_from_devt(dev_t devt, u64 *idx_mask)
{
        for (int type = RDMA_UCAP_FIRST; type < RDMA_UCAP_MAX; type++) {
                if (ucaps_list[type] && ucaps_list[type]->dev.devt == devt) {
                        *idx_mask |= 1 << type;
                        return 0;
                }
        }

        return -EINVAL;
}

static int get_devt_from_fd(unsigned int fd, dev_t *ret_dev)
{
        struct file *file;

        file = fget(fd);
        if (!file)
                return -EBADF;

        *ret_dev = file_inode(file)->i_rdev;
        fput(file);
        return 0;
}

/**
 * ib_ucaps_init - Initialization required before ucap creation.
 *
 * Return: 0 on success, or a negative errno value on failure
 */
static int ib_ucaps_init(void)
{
        int ret = 0;

        if (ucaps_class_is_registered)
                return ret;

        ret = class_register(&ucaps_class);
        if (ret)
                return ret;

        ret = alloc_chrdev_region(&ucaps_base_dev, 0, RDMA_UCAP_MAX,
                                  ucaps_class.name);
        if (ret < 0) {
                class_unregister(&ucaps_class);
                return ret;
        }

        ucaps_class_is_registered = true;

        return 0;
}

static void ucap_dev_release(struct device *device)
{
        struct ib_ucap *ucap = container_of(device, struct ib_ucap, dev);

        kfree(ucap);
}

/**
 * ib_create_ucap - Add a ucap character device
 * @type: UCAP type
 *
 * Creates a ucap character device in the /dev/infiniband directory. By default,
 * the device has root-only read-write access.
 *
 * A driver may call this multiple times with the same UCAP type. A reference
 * count tracks creations and deletions.
 *
 * Return: 0 on success, or a negative errno value on failure
 */
int ib_create_ucap(enum rdma_user_cap type)
{
        struct ib_ucap *ucap;
        int ret;

        if (type >= RDMA_UCAP_MAX)
                return -EINVAL;

        mutex_lock(&ucaps_mutex);
        ret = ib_ucaps_init();
        if (ret)
                goto unlock;

        ucap = ucaps_list[type];
        if (ucap) {
                kref_get(&ucap->ref);
                mutex_unlock(&ucaps_mutex);
                return 0;
        }

        ucap = kzalloc_obj(*ucap);
        if (!ucap) {
                ret = -ENOMEM;
                goto unlock;
        }

        device_initialize(&ucap->dev);
        ucap->dev.class = &ucaps_class;
        ucap->dev.devt = MKDEV(MAJOR(ucaps_base_dev), type);
        ucap->dev.release = ucap_dev_release;
        ret = dev_set_name(&ucap->dev, "%s", ucap_names[type]);
        if (ret)
                goto err_device;

        cdev_init(&ucap->cdev, &ucaps_cdev_fops);
        ucap->cdev.owner = THIS_MODULE;

        ret = cdev_device_add(&ucap->cdev, &ucap->dev);
        if (ret)
                goto err_device;

        kref_init(&ucap->ref);
        ucaps_list[type] = ucap;
        mutex_unlock(&ucaps_mutex);

        return 0;

err_device:
        put_device(&ucap->dev);
unlock:
        mutex_unlock(&ucaps_mutex);
        return ret;
}
EXPORT_SYMBOL(ib_create_ucap);

static void ib_release_ucap(struct kref *ref)
{
        struct ib_ucap *ucap = container_of(ref, struct ib_ucap, ref);
        enum rdma_user_cap type;

        for (type = RDMA_UCAP_FIRST; type < RDMA_UCAP_MAX; type++) {
                if (ucaps_list[type] == ucap)
                        break;
        }
        WARN_ON(type == RDMA_UCAP_MAX);

        ucaps_list[type] = NULL;
        cdev_device_del(&ucap->cdev, &ucap->dev);
        put_device(&ucap->dev);
}

/**
 * ib_remove_ucap - Remove a ucap character device
 * @type: User cap type
 *
 * Removes the ucap character device according to type. The device is completely
 * removed from the filesystem when its reference count reaches 0.
 */
void ib_remove_ucap(enum rdma_user_cap type)
{
        struct ib_ucap *ucap;

        mutex_lock(&ucaps_mutex);
        ucap = ucaps_list[type];
        if (WARN_ON(!ucap))
                goto end;

        kref_put(&ucap->ref, ib_release_ucap);
end:
        mutex_unlock(&ucaps_mutex);
}
EXPORT_SYMBOL(ib_remove_ucap);

/**
 * ib_get_ucaps - Get bitmask of ucap types from file descriptors
 * @fds: Array of file descriptors
 * @fd_count: Number of file descriptors in the array
 * @idx_mask: Bitmask to be updated based on the ucaps in the fd list
 *
 * Given an array of file descriptors, this function returns a bitmask of
 * the ucaps where a bit is set if an FD for that ucap type was in the array.
 *
 * Return: 0 on success, or a negative errno value on failure
 */
int ib_get_ucaps(int *fds, int fd_count, uint64_t *idx_mask)
{
        int ret = 0;
        dev_t dev;

        *idx_mask = 0;
        mutex_lock(&ucaps_mutex);
        for (int i = 0; i < fd_count; i++) {
                ret = get_devt_from_fd(fds[i], &dev);
                if (ret)
                        goto end;

                ret = get_ucap_from_devt(dev, idx_mask);
                if (ret)
                        goto end;
        }

end:
        mutex_unlock(&ucaps_mutex);
        return ret;
}