// SPDX-License-Identifier: GPL-2.0-or-later
/*
 *      Devmem TCP
 *
 *      Authors:	Mina Almasry <almasrymina@google.com>
 *			Willem de Bruijn <willemdebruijn.kernel@gmail.com>
 *			Kaiyuan Zhang <kaiyuanz@google.com
 */

#include <linux/dma-buf.h>
#include <linux/genalloc.h>
#include <linux/mm.h>
#include <linux/netdevice.h>
#include <linux/types.h>
#include <net/netdev_queues.h>
#include <net/netdev_rx_queue.h>
#include <net/page_pool/helpers.h>
#include <net/page_pool/memory_provider.h>
#include <net/sock.h>
#include <net/tcp.h>
#include <trace/events/page_pool.h>

#include "devmem.h"
#include "mp_dmabuf_devmem.h"
#include "page_pool_priv.h"

/* Device memory support */

static DEFINE_XARRAY_FLAGS(net_devmem_dmabuf_bindings, XA_FLAGS_ALLOC1);

static const struct memory_provider_ops dmabuf_devmem_ops;

static void net_devmem_dmabuf_free_chunk_owner(struct gen_pool *genpool,
					       struct gen_pool_chunk *chunk,
					       void *not_used)
{
	struct dmabuf_genpool_chunk_owner *owner = chunk->owner;

	kvfree(owner->area.niovs);
	kfree(owner);
}

static dma_addr_t net_devmem_get_dma_addr(const struct net_iov *niov)
{
	struct dmabuf_genpool_chunk_owner *owner;

	owner = net_devmem_iov_to_chunk_owner(niov);
	return owner->base_dma_addr +
	       ((dma_addr_t)net_iov_idx(niov) << owner->binding->niov_shift);
}

static void net_devmem_dmabuf_binding_release(struct percpu_ref *ref)
{
	struct net_devmem_dmabuf_binding *binding =
		container_of(ref, struct net_devmem_dmabuf_binding, ref);

	INIT_WORK(&binding->unbind_w, __net_devmem_dmabuf_binding_free);
	schedule_work(&binding->unbind_w);
}

void __net_devmem_dmabuf_binding_free(struct work_struct *wq)
{
	struct net_devmem_dmabuf_binding *binding = container_of(wq, typeof(*binding), unbind_w);

	size_t size, avail;

	gen_pool_for_each_chunk(binding->chunk_pool,
				net_devmem_dmabuf_free_chunk_owner, NULL);

	size = gen_pool_size(binding->chunk_pool);
	avail = gen_pool_avail(binding->chunk_pool);

	if (!WARN(size != avail, "can't destroy genpool. size=%zu, avail=%zu",
		  size, avail))
		gen_pool_destroy(binding->chunk_pool);

	dma_buf_unmap_attachment_unlocked(binding->attachment, binding->sgt,
					  binding->direction);
	dma_buf_detach(binding->dmabuf, binding->attachment);
	dma_buf_put(binding->dmabuf);
	xa_destroy(&binding->bound_rxqs);
	percpu_ref_exit(&binding->ref);
	kvfree(binding->tx_vec);
	kfree(binding);
}

struct net_iov *
net_devmem_alloc_dmabuf(struct net_devmem_dmabuf_binding *binding)
{
	struct dmabuf_genpool_chunk_owner *owner;
	unsigned long dma_addr;
	struct net_iov *niov;
	ssize_t offset;
	ssize_t index;

	dma_addr = gen_pool_alloc_owner(binding->chunk_pool,
					1UL << binding->niov_shift,
					(void **)&owner);
	if (!dma_addr)
		return NULL;

	offset = dma_addr - owner->base_dma_addr;
	index = offset >> binding->niov_shift;
	niov = &owner->area.niovs[index];

	niov->desc.pp_magic = 0;
	niov->desc.pp = NULL;
	atomic_long_set(&niov->desc.pp_ref_count, 0);

	return niov;
}

void net_devmem_free_dmabuf(struct net_iov *niov)
{
	struct net_devmem_dmabuf_binding *binding = net_devmem_iov_binding(niov);
	unsigned long dma_addr = net_devmem_get_dma_addr(niov);
	size_t niov_size = 1UL << binding->niov_shift;

	if (WARN_ON(!gen_pool_has_addr(binding->chunk_pool, dma_addr,
				       niov_size)))
		return;

	gen_pool_free(binding->chunk_pool, dma_addr, niov_size);
}

void net_devmem_unbind_dmabuf(struct net_devmem_dmabuf_binding *binding)
{
	struct netdev_rx_queue *rxq;
	unsigned long xa_idx;
	unsigned int rxq_idx;

	xa_erase(&net_devmem_dmabuf_bindings, binding->id);

	/* Ensure no tx net_devmem_lookup_dmabuf() are in flight after the
	 * erase.
	 */
	synchronize_net();

	if (binding->list.next)
		list_del(&binding->list);

	xa_for_each(&binding->bound_rxqs, xa_idx, rxq) {
		const struct pp_memory_provider_params mp_params = {
			.mp_priv	= binding,
			.mp_ops		= &dmabuf_devmem_ops,
		};

		rxq_idx = get_netdev_rx_queue_index(rxq);

		netif_mp_close_rxq(binding->dev, rxq_idx, &mp_params);
	}

	percpu_ref_kill(&binding->ref);
}

int net_devmem_bind_dmabuf_to_queue(struct net_device *dev, u32 rxq_idx,
				    struct net_devmem_dmabuf_binding *binding,
				    struct netlink_ext_ack *extack)
{
	struct pp_memory_provider_params mp_params = {
		.mp_priv	= binding,
		.mp_ops		= &dmabuf_devmem_ops,
	};
	struct netdev_rx_queue *rxq;
	u32 xa_idx;
	int err;

	if (binding->niov_shift != PAGE_SHIFT)
		mp_params.rx_page_size = 1U << binding->niov_shift;

	err = netif_mp_open_rxq(dev, rxq_idx, &mp_params, extack);
	if (err)
		return err;

	rxq = __netif_get_rx_queue(dev, rxq_idx);
	err = xa_alloc(&binding->bound_rxqs, &xa_idx, rxq, xa_limit_32b,
		       GFP_KERNEL);
	if (err)
		goto err_close_rxq;

	return 0;

err_close_rxq:
	netif_mp_close_rxq(dev, rxq_idx, &mp_params);
	return err;
}

struct net_devmem_dmabuf_binding *
net_devmem_bind_dmabuf(struct net_device *dev, void *vdev,
		       struct device *dma_dev,
		       enum dma_data_direction direction,
		       unsigned int dmabuf_fd, unsigned int niov_shift,
		       struct netdev_nl_sock *priv,
		       struct netlink_ext_ack *extack)
{
	struct net_devmem_dmabuf_binding *binding;
	size_t niov_size = 1UL << niov_shift;
	static u32 id_alloc_next;
	struct scatterlist *sg;
	struct dma_buf *dmabuf;
	unsigned int sg_idx, i;
	unsigned long virtual;
	int err;

	if (!dma_dev) {
		NL_SET_ERR_MSG(extack, "Device doesn't support DMA");
		return ERR_PTR(-EOPNOTSUPP);
	}

	dmabuf = dma_buf_get(dmabuf_fd);
	if (IS_ERR(dmabuf))
		return ERR_CAST(dmabuf);

	binding = kzalloc_node(sizeof(*binding), GFP_KERNEL,
			       dev_to_node(&dev->dev));
	if (!binding) {
		err = -ENOMEM;
		goto err_put_dmabuf;
	}

	binding->dev = dev;
	binding->vdev = vdev;
	binding->niov_shift = niov_shift;
	xa_init_flags(&binding->bound_rxqs, XA_FLAGS_ALLOC);

	err = percpu_ref_init(&binding->ref,
			      net_devmem_dmabuf_binding_release,
			      0, GFP_KERNEL);
	if (err < 0)
		goto err_free_binding;

	mutex_init(&binding->lock);

	binding->dmabuf = dmabuf;
	binding->direction = direction;

	binding->attachment = dma_buf_attach(binding->dmabuf, dma_dev);
	if (IS_ERR(binding->attachment)) {
		err = PTR_ERR(binding->attachment);
		NL_SET_ERR_MSG(extack, "Failed to bind dmabuf to device");
		goto err_exit_ref;
	}

	binding->sgt = dma_buf_map_attachment_unlocked(binding->attachment,
						       direction);
	if (IS_ERR(binding->sgt)) {
		err = PTR_ERR(binding->sgt);
		NL_SET_ERR_MSG(extack, "Failed to map dmabuf attachment");
		goto err_detach;
	}

	if (direction == DMA_TO_DEVICE) {
		if (!IS_ALIGNED(dmabuf->size, PAGE_SIZE)) {
			err = -EINVAL;
			NL_SET_ERR_MSG(extack, "TX dma-buf size must be a multiple of PAGE_SIZE");
			goto err_unmap;
		}
		binding->tx_vec = kvmalloc_objs(struct net_iov *,
						dmabuf->size / PAGE_SIZE);
		if (!binding->tx_vec) {
			err = -ENOMEM;
			goto err_unmap;
		}
	}

	binding->chunk_pool = gen_pool_create(niov_shift,
					      dev_to_node(&dev->dev));
	if (!binding->chunk_pool) {
		err = -ENOMEM;
		goto err_tx_vec;
	}

	virtual = 0;
	for_each_sgtable_dma_sg(binding->sgt, sg, sg_idx) {
		dma_addr_t dma_addr = sg_dma_address(sg);
		struct dmabuf_genpool_chunk_owner *owner;
		size_t len = sg_dma_len(sg);
		struct net_iov *niov;

		if (!IS_ALIGNED(dma_addr, niov_size) ||
		    !IS_ALIGNED(len, niov_size)) {
			err = -EINVAL;
			NL_SET_ERR_MSG_FMT(extack,
					   "dmabuf sg entry (addr=%pad, len=%zu) not aligned to niov size %zu",
					   &dma_addr, len, niov_size);
			goto err_free_chunks;
		}

		owner = kzalloc_node(sizeof(*owner), GFP_KERNEL,
				     dev_to_node(&dev->dev));
		if (!owner) {
			err = -ENOMEM;
			goto err_free_chunks;
		}

		owner->area.base_virtual = virtual;
		owner->base_dma_addr = dma_addr;
		owner->area.num_niovs = len >> niov_shift;
		owner->binding = binding;

		err = gen_pool_add_owner(binding->chunk_pool, dma_addr,
					 dma_addr, len, dev_to_node(&dev->dev),
					 owner);
		if (err) {
			kfree(owner);
			err = -EINVAL;
			goto err_free_chunks;
		}

		owner->area.niovs = kvmalloc_objs(*owner->area.niovs,
						  owner->area.num_niovs);
		if (!owner->area.niovs) {
			err = -ENOMEM;
			goto err_free_chunks;
		}

		for (i = 0; i < owner->area.num_niovs; i++) {
			niov = &owner->area.niovs[i];
			net_iov_init(niov, &owner->area, NET_IOV_DMABUF);
			page_pool_set_dma_addr_netmem(net_iov_to_netmem(niov),
						      net_devmem_get_dma_addr(niov));
			if (direction == DMA_TO_DEVICE)
				binding->tx_vec[owner->area.base_virtual / PAGE_SIZE + i] = niov;
		}

		virtual += len;
	}

	err = xa_alloc_cyclic(&net_devmem_dmabuf_bindings, &binding->id,
			      binding, xa_limit_32b, &id_alloc_next,
			      GFP_KERNEL);
	if (err < 0)
		goto err_free_chunks;

	list_add(&binding->list, &priv->bindings);

	return binding;

err_free_chunks:
	gen_pool_for_each_chunk(binding->chunk_pool,
				net_devmem_dmabuf_free_chunk_owner, NULL);
	gen_pool_destroy(binding->chunk_pool);
err_tx_vec:
	kvfree(binding->tx_vec);
err_unmap:
	dma_buf_unmap_attachment_unlocked(binding->attachment, binding->sgt,
					  direction);
err_detach:
	dma_buf_detach(dmabuf, binding->attachment);
err_exit_ref:
	percpu_ref_exit(&binding->ref);
err_free_binding:
	kfree(binding);
err_put_dmabuf:
	dma_buf_put(dmabuf);
	return ERR_PTR(err);
}

struct net_devmem_dmabuf_binding *net_devmem_lookup_dmabuf(u32 id)
{
	struct net_devmem_dmabuf_binding *binding;

	rcu_read_lock();
	binding = xa_load(&net_devmem_dmabuf_bindings, id);
	if (binding) {
		if (!net_devmem_dmabuf_binding_get(binding))
			binding = NULL;
	}
	rcu_read_unlock();

	return binding;
}

void net_devmem_get_net_iov(struct net_iov *niov)
{
	net_devmem_dmabuf_binding_get(net_devmem_iov_binding(niov));
}

void net_devmem_put_net_iov(struct net_iov *niov)
{
	net_devmem_dmabuf_binding_put(net_devmem_iov_binding(niov));
}

struct net_devmem_dmabuf_binding *net_devmem_get_binding(struct sock *sk,
							 unsigned int dmabuf_id)
{
	struct net_devmem_dmabuf_binding *binding;
	struct net_device *dst_dev;
	struct dst_entry *dst;
	int err = 0;

	binding = net_devmem_lookup_dmabuf(dmabuf_id);
	if (!binding || !binding->tx_vec) {
		err = -EINVAL;
		goto out_err;
	}

	rcu_read_lock();
	dst = __sk_dst_get(sk);
	/* If dst is NULL (route expired), attempt to rebuild it. */
	if (unlikely(!dst)) {
		if (inet_csk(sk)->icsk_af_ops->rebuild_header(sk)) {
			err = -EHOSTUNREACH;
			goto out_unlock;
		}
		dst = __sk_dst_get(sk);
		if (unlikely(!dst)) {
			err = -ENODEV;
			goto out_unlock;
		}
	}

	/* The dma-addrs in this binding are only reachable to the corresponding
	 * net_device.
	 */
	dst_dev = dst_dev_rcu(dst);
	if (unlikely(!dst_dev) ||
	    unlikely(dst_dev != READ_ONCE(binding->dev) &&
		     dst_dev != READ_ONCE(binding->vdev))) {
		err = -ENODEV;
		goto out_unlock;
	}

	rcu_read_unlock();
	return binding;

out_unlock:
	rcu_read_unlock();
out_err:
	if (binding)
		net_devmem_dmabuf_binding_put(binding);

	return ERR_PTR(err);
}

struct net_iov *
net_devmem_get_niov_at(struct net_devmem_dmabuf_binding *binding,
		       size_t virt_addr, size_t *off, size_t *size)
{
	if (virt_addr >= binding->dmabuf->size)
		return NULL;

	*off = virt_addr % PAGE_SIZE;
	*size = PAGE_SIZE - *off;

	return binding->tx_vec[virt_addr / PAGE_SIZE];
}

/*** "Dmabuf devmem memory provider" ***/

int mp_dmabuf_devmem_init(struct page_pool *pool)
{
	struct net_devmem_dmabuf_binding *binding = pool->mp_priv;

	if (!binding)
		return -EINVAL;

	/* dma-buf dma addresses do not need and should not be used with
	 * dma_sync_for_cpu/device. Force disable dma_sync.
	 */
	pool->dma_sync = false;
	pool->dma_sync_for_cpu = false;

	if (pool->p.order != binding->niov_shift - PAGE_SHIFT)
		return -E2BIG;

	net_devmem_dmabuf_binding_get(binding);
	return 0;
}

netmem_ref mp_dmabuf_devmem_alloc_netmems(struct page_pool *pool, gfp_t gfp)
{
	struct net_devmem_dmabuf_binding *binding = pool->mp_priv;
	struct net_iov *niov;
	netmem_ref netmem;

	niov = net_devmem_alloc_dmabuf(binding);
	if (!niov)
		return 0;

	netmem = net_iov_to_netmem(niov);

	page_pool_set_pp_info(pool, netmem);

	pool->pages_state_hold_cnt++;
	trace_page_pool_state_hold(pool, netmem, pool->pages_state_hold_cnt);
	return netmem;
}

void mp_dmabuf_devmem_destroy(struct page_pool *pool)
{
	struct net_devmem_dmabuf_binding *binding = pool->mp_priv;

	net_devmem_dmabuf_binding_put(binding);
}

bool mp_dmabuf_devmem_release_page(struct page_pool *pool, netmem_ref netmem)
{
	long refcount = atomic_long_read(netmem_get_pp_ref_count_ref(netmem));

	if (WARN_ON_ONCE(!netmem_is_net_iov(netmem)))
		return false;

	if (WARN_ON_ONCE(refcount != 1))
		return false;

	page_pool_clear_pp_info(netmem);

	net_devmem_free_dmabuf(netmem_to_net_iov(netmem));

	/* We don't want the page pool put_page()ing our net_iovs. */
	return false;
}

static int mp_dmabuf_devmem_nl_fill(void *mp_priv, struct sk_buff *rsp,
				    struct netdev_rx_queue *rxq)
{
	const struct net_devmem_dmabuf_binding *binding = mp_priv;
	int type = rxq ? NETDEV_A_QUEUE_DMABUF : NETDEV_A_PAGE_POOL_DMABUF;

	return nla_put_u32(rsp, type, binding->id);
}

static void mp_dmabuf_devmem_uninstall(void *mp_priv,
				       struct netdev_rx_queue *rxq)
{
	struct net_devmem_dmabuf_binding *binding = mp_priv;
	struct netdev_rx_queue *bound_rxq;
	unsigned long xa_idx;

	xa_for_each(&binding->bound_rxqs, xa_idx, bound_rxq) {
		if (bound_rxq == rxq) {
			xa_erase(&binding->bound_rxqs, xa_idx);
			if (xa_empty(&binding->bound_rxqs)) {
				mutex_lock(&binding->lock);
				ASSERT_EXCLUSIVE_WRITER(binding->dev);
				WRITE_ONCE(binding->dev, NULL);
				mutex_unlock(&binding->lock);
			}
			break;
		}
	}
}

static const struct memory_provider_ops dmabuf_devmem_ops = {
	.init			= mp_dmabuf_devmem_init,
	.destroy		= mp_dmabuf_devmem_destroy,
	.alloc_netmems		= mp_dmabuf_devmem_alloc_netmems,
	.release_netmem		= mp_dmabuf_devmem_release_page,
	.nl_fill		= mp_dmabuf_devmem_nl_fill,
	.uninstall		= mp_dmabuf_devmem_uninstall,
};