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

#include <linux/dma-buf-mapping.h>
#include <linux/pci-p2pdma.h>
#include <linux/dma-resv.h>
#include <rdma/uverbs_std_types.h>
#include "rdma_core.h"
#include "uverbs.h"

MODULE_IMPORT_NS("DMA_BUF");

static int uverbs_dmabuf_attach(struct dma_buf *dmabuf,
                                struct dma_buf_attachment *attachment)
{
        if (!attachment->peer2peer)
                return -EOPNOTSUPP;

        return 0;
}

static struct sg_table *
uverbs_dmabuf_map(struct dma_buf_attachment *attachment,
                  enum dma_data_direction dir)
{
        struct ib_uverbs_dmabuf_file *priv = attachment->dmabuf->priv;
        struct sg_table *ret;

        dma_resv_assert_held(priv->dmabuf->resv);

        if (priv->revoked)
                return ERR_PTR(-ENODEV);

        ret = dma_buf_phys_vec_to_sgt(attachment, priv->provider,
                                      &priv->phys_vec, 1, priv->phys_vec.len,
                                      dir);
        if (IS_ERR(ret))
                return ret;

        kref_get(&priv->kref);
        return ret;
}

static void uverbs_dmabuf_unmap(struct dma_buf_attachment *attachment,
                                struct sg_table *sgt,
                                enum dma_data_direction dir)
{
        struct ib_uverbs_dmabuf_file *priv = attachment->dmabuf->priv;

        dma_resv_assert_held(priv->dmabuf->resv);
        dma_buf_free_sgt(attachment, sgt, dir);
        kref_put(&priv->kref, ib_uverbs_dmabuf_done);
}

static int uverbs_dmabuf_pin(struct dma_buf_attachment *attach)
{
        return -EOPNOTSUPP;
}

static void uverbs_dmabuf_unpin(struct dma_buf_attachment *attach)
{
}

static void uverbs_dmabuf_release(struct dma_buf *dmabuf)
{
        struct ib_uverbs_dmabuf_file *priv = dmabuf->priv;

        /*
         * This can only happen if the fput came from alloc_abort_fd_uobject()
         */
        if (!priv->uobj.context)
                return;

        uverbs_uobject_release(&priv->uobj);
}

static const struct dma_buf_ops uverbs_dmabuf_ops = {
        .attach = uverbs_dmabuf_attach,
        .map_dma_buf = uverbs_dmabuf_map,
        .unmap_dma_buf = uverbs_dmabuf_unmap,
        .pin = uverbs_dmabuf_pin,
        .unpin = uverbs_dmabuf_unpin,
        .release = uverbs_dmabuf_release,
};

static int UVERBS_HANDLER(UVERBS_METHOD_DMABUF_ALLOC)(
        struct uverbs_attr_bundle *attrs)
{
        struct ib_uobject *uobj =
                uverbs_attr_get(attrs, UVERBS_ATTR_ALLOC_DMABUF_HANDLE)
                        ->obj_attr.uobject;
        struct ib_uverbs_dmabuf_file *uverbs_dmabuf =
                container_of(uobj, struct ib_uverbs_dmabuf_file, uobj);
        struct ib_device *ib_dev = attrs->context->device;
        struct rdma_user_mmap_entry *mmap_entry;
        DEFINE_DMA_BUF_EXPORT_INFO(exp_info);
        off_t pg_off;
        int ret;

        ret = uverbs_get_const(&pg_off, attrs, UVERBS_ATTR_ALLOC_DMABUF_PGOFF);
        if (ret)
                return ret;

        mmap_entry = ib_dev->ops.pgoff_to_mmap_entry(attrs->context, pg_off);
        if (!mmap_entry)
                return -EINVAL;

        ret = ib_dev->ops.mmap_get_pfns(mmap_entry, &uverbs_dmabuf->phys_vec,
                                        &uverbs_dmabuf->provider);
        if (ret)
                goto err;

        exp_info.ops = &uverbs_dmabuf_ops;
        exp_info.size = uverbs_dmabuf->phys_vec.len;
        exp_info.flags = O_CLOEXEC;
        exp_info.priv = uverbs_dmabuf;

        uverbs_dmabuf->dmabuf = dma_buf_export(&exp_info);
        if (IS_ERR(uverbs_dmabuf->dmabuf)) {
                ret = PTR_ERR(uverbs_dmabuf->dmabuf);
                goto err;
        }

        kref_init(&uverbs_dmabuf->kref);
        init_completion(&uverbs_dmabuf->comp);
        INIT_LIST_HEAD(&uverbs_dmabuf->dmabufs_elm);
        mutex_lock(&mmap_entry->dmabufs_lock);
        if (mmap_entry->driver_removed)
                ret = -EIO;
        else
                list_add_tail(&uverbs_dmabuf->dmabufs_elm, &mmap_entry->dmabufs);
        mutex_unlock(&mmap_entry->dmabufs_lock);
        if (ret)
                goto err_revoked;

        uobj->object = uverbs_dmabuf->dmabuf->file;
        uverbs_dmabuf->mmap_entry = mmap_entry;
        uverbs_finalize_uobj_create(attrs, UVERBS_ATTR_ALLOC_DMABUF_HANDLE);
        return 0;

err_revoked:
        dma_buf_put(uverbs_dmabuf->dmabuf);
err:
        rdma_user_mmap_entry_put(mmap_entry);
        return ret;
}

DECLARE_UVERBS_NAMED_METHOD(
        UVERBS_METHOD_DMABUF_ALLOC,
        UVERBS_ATTR_FD(UVERBS_ATTR_ALLOC_DMABUF_HANDLE,
                       UVERBS_OBJECT_DMABUF,
                       UVERBS_ACCESS_NEW,
                       UA_MANDATORY),
        UVERBS_ATTR_PTR_IN(UVERBS_ATTR_ALLOC_DMABUF_PGOFF,
                           UVERBS_ATTR_TYPE(u64),
                           UA_MANDATORY));

static void uverbs_dmabuf_fd_destroy_uobj(struct ib_uobject *uobj,
                                          enum rdma_remove_reason why)
{
        struct ib_uverbs_dmabuf_file *uverbs_dmabuf =
                container_of(uobj, struct ib_uverbs_dmabuf_file, uobj);
        bool wait_for_comp = false;

        mutex_lock(&uverbs_dmabuf->mmap_entry->dmabufs_lock);
        dma_resv_lock(uverbs_dmabuf->dmabuf->resv, NULL);
        if (!uverbs_dmabuf->revoked) {
                uverbs_dmabuf->revoked = true;
                list_del(&uverbs_dmabuf->dmabufs_elm);
                dma_buf_move_notify(uverbs_dmabuf->dmabuf);
                dma_resv_wait_timeout(uverbs_dmabuf->dmabuf->resv,
                                      DMA_RESV_USAGE_BOOKKEEP, false,
                                      MAX_SCHEDULE_TIMEOUT);
                wait_for_comp = true;
        }
        dma_resv_unlock(uverbs_dmabuf->dmabuf->resv);
        if (wait_for_comp) {
                kref_put(&uverbs_dmabuf->kref, ib_uverbs_dmabuf_done);
                /* Let's wait till all DMA unmap are completed. */
                wait_for_completion(&uverbs_dmabuf->comp);
        }
        mutex_unlock(&uverbs_dmabuf->mmap_entry->dmabufs_lock);

        /* Matches the get done as part of pgoff_to_mmap_entry() */
        rdma_user_mmap_entry_put(uverbs_dmabuf->mmap_entry);
}

DECLARE_UVERBS_NAMED_OBJECT(
        UVERBS_OBJECT_DMABUF,
        UVERBS_TYPE_ALLOC_FD(sizeof(struct ib_uverbs_dmabuf_file),
                             uverbs_dmabuf_fd_destroy_uobj,
                             NULL, NULL, O_RDONLY),
                             &UVERBS_METHOD(UVERBS_METHOD_DMABUF_ALLOC));

const struct uapi_definition uverbs_def_obj_dmabuf[] = {
        UAPI_DEF_CHAIN_OBJ_TREE_NAMED(UVERBS_OBJECT_DMABUF),
                                      UAPI_DEF_OBJ_NEEDS_FN(mmap_get_pfns),
                                      UAPI_DEF_OBJ_NEEDS_FN(pgoff_to_mmap_entry),
        {}
};