// SPDX-License-Identifier: GPL-2.0-only
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/netfilter.h>
#include <linux/rhashtable.h>
#include <linux/netdevice.h>
#include <net/flow_offload.h>
#include <net/netfilter/nf_flow_table.h>

struct flow_offload_xdp_ft {
        struct list_head head;
        struct nf_flowtable *ft;
        struct rcu_head rcuhead;
};

struct flow_offload_xdp {
        struct hlist_node hnode;
        unsigned long net_device_addr;
        struct list_head head;
};

#define NF_XDP_HT_BITS  4
static DEFINE_HASHTABLE(nf_xdp_hashtable, NF_XDP_HT_BITS);
static DEFINE_MUTEX(nf_xdp_hashtable_lock);

/* caller must hold rcu read lock */
struct nf_flowtable *nf_flowtable_by_dev(const struct net_device *dev)
{
        unsigned long key = (unsigned long)dev;
        struct flow_offload_xdp *iter;

        hash_for_each_possible_rcu(nf_xdp_hashtable, iter, hnode, key) {
                if (key == iter->net_device_addr) {
                        struct flow_offload_xdp_ft *ft_elem;

                        /* The user is supposed to insert a given net_device
                         * just into a single nf_flowtable so we always return
                         * the first element here.
                         */
                        ft_elem = list_first_or_null_rcu(&iter->head,
                                                         struct flow_offload_xdp_ft,
                                                         head);
                        return ft_elem ? ft_elem->ft : NULL;
                }
        }

        return NULL;
}

static int nf_flowtable_by_dev_insert(struct nf_flowtable *ft,
                                      const struct net_device *dev)
{
        struct flow_offload_xdp *iter, *elem = NULL;
        unsigned long key = (unsigned long)dev;
        struct flow_offload_xdp_ft *ft_elem;

        ft_elem = kzalloc_obj(*ft_elem, GFP_KERNEL_ACCOUNT);
        if (!ft_elem)
                return -ENOMEM;

        ft_elem->ft = ft;

        mutex_lock(&nf_xdp_hashtable_lock);

        hash_for_each_possible(nf_xdp_hashtable, iter, hnode, key) {
                if (key == iter->net_device_addr) {
                        elem = iter;
                        break;
                }
        }

        if (!elem) {
                elem = kzalloc_obj(*elem, GFP_KERNEL_ACCOUNT);
                if (!elem)
                        goto err_unlock;

                elem->net_device_addr = key;
                INIT_LIST_HEAD(&elem->head);
                hash_add_rcu(nf_xdp_hashtable, &elem->hnode, key);
        }
        list_add_tail_rcu(&ft_elem->head, &elem->head);

        mutex_unlock(&nf_xdp_hashtable_lock);

        return 0;

err_unlock:
        mutex_unlock(&nf_xdp_hashtable_lock);
        kfree(ft_elem);

        return -ENOMEM;
}

static void nf_flowtable_by_dev_remove(struct nf_flowtable *ft,
                                       const struct net_device *dev)
{
        struct flow_offload_xdp *iter, *elem = NULL;
        unsigned long key = (unsigned long)dev;

        mutex_lock(&nf_xdp_hashtable_lock);

        hash_for_each_possible(nf_xdp_hashtable, iter, hnode, key) {
                if (key == iter->net_device_addr) {
                        elem = iter;
                        break;
                }
        }

        if (elem) {
                struct flow_offload_xdp_ft *ft_elem, *ft_next;

                list_for_each_entry_safe(ft_elem, ft_next, &elem->head, head) {
                        if (ft_elem->ft == ft) {
                                list_del_rcu(&ft_elem->head);
                                kfree_rcu(ft_elem, rcuhead);
                        }
                }

                if (list_empty(&elem->head))
                        hash_del_rcu(&elem->hnode);
                else
                        elem = NULL;
        }

        mutex_unlock(&nf_xdp_hashtable_lock);

        if (elem) {
                synchronize_rcu();
                kfree(elem);
        }
}

int nf_flow_offload_xdp_setup(struct nf_flowtable *flowtable,
                              struct net_device *dev,
                              enum flow_block_command cmd)
{
        switch (cmd) {
        case FLOW_BLOCK_BIND:
                return nf_flowtable_by_dev_insert(flowtable, dev);
        case FLOW_BLOCK_UNBIND:
                nf_flowtable_by_dev_remove(flowtable, dev);
                return 0;
        }

        WARN_ON_ONCE(1);
        return 0;
}