path: root/net/netfilter/nft_set_rbtree.c

// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (c) 2008-2009 Patrick McHardy <kaber@trash.net>
 *
 * Development of this code funded by Astaro AG (http://www.astaro.com/)
 */

#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/list.h>
#include <linux/rbtree.h>
#include <linux/netlink.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nf_tables.h>
#include <net/netfilter/nf_tables_core.h>

struct nft_rbtree {
	struct rb_root		root;
	rwlock_t		lock;
	seqcount_rwlock_t	count;
	unsigned long		last_gc;
};

struct nft_rbtree_elem {
	struct nft_elem_priv	priv;
	struct rb_node		node;
	struct nft_set_ext	ext;
};

static bool nft_rbtree_interval_end(const struct nft_rbtree_elem *rbe)
{
	return nft_set_ext_exists(&rbe->ext, NFT_SET_EXT_FLAGS) &&
	       (*nft_set_ext_flags(&rbe->ext) & NFT_SET_ELEM_INTERVAL_END);
}

static bool nft_rbtree_interval_start(const struct nft_rbtree_elem *rbe)
{
	return !nft_rbtree_interval_end(rbe);
}

static int nft_rbtree_cmp(const struct nft_set *set,
			  const struct nft_rbtree_elem *e1,
			  const struct nft_rbtree_elem *e2)
{
	return memcmp(nft_set_ext_key(&e1->ext), nft_set_ext_key(&e2->ext),
		      set->klen);
}

static bool nft_rbtree_elem_expired(const struct nft_rbtree_elem *rbe)
{
	return nft_set_elem_expired(&rbe->ext);
}

static const struct nft_set_ext *
__nft_rbtree_lookup(const struct net *net, const struct nft_set *set,
		    const u32 *key, unsigned int seq)
{
	struct nft_rbtree *priv = nft_set_priv(set);
	const struct nft_rbtree_elem *rbe, *interval = NULL;
	u8 genmask = nft_genmask_cur(net);
	const struct rb_node *parent;
	int d;

	parent = rcu_dereference_raw(priv->root.rb_node);
	while (parent != NULL) {
		if (read_seqcount_retry(&priv->count, seq))
			return NULL;

		rbe = rb_entry(parent, struct nft_rbtree_elem, node);

		d = memcmp(nft_set_ext_key(&rbe->ext), key, set->klen);
		if (d < 0) {
			parent = rcu_dereference_raw(parent->rb_left);
			if (interval &&
			    !nft_rbtree_cmp(set, rbe, interval) &&
			    nft_rbtree_interval_end(rbe) &&
			    nft_rbtree_interval_start(interval))
				continue;
			if (nft_set_elem_active(&rbe->ext, genmask) &&
			    !nft_rbtree_elem_expired(rbe))
				interval = rbe;
		} else if (d > 0)
			parent = rcu_dereference_raw(parent->rb_right);
		else {
			if (!nft_set_elem_active(&rbe->ext, genmask)) {
				parent = rcu_dereference_raw(parent->rb_left);
				continue;
			}

			if (nft_rbtree_elem_expired(rbe))
				return NULL;

			if (nft_rbtree_interval_end(rbe)) {
				if (nft_set_is_anonymous(set))
					return NULL;
				parent = rcu_dereference_raw(parent->rb_left);
				interval = NULL;
				continue;
			}

			return &rbe->ext;
		}
	}

	if (set->flags & NFT_SET_INTERVAL && interval != NULL &&
	    nft_rbtree_interval_start(interval))
		return &interval->ext;

	return NULL;
}

INDIRECT_CALLABLE_SCOPE
const struct nft_set_ext *
nft_rbtree_lookup(const struct net *net, const struct nft_set *set,
		  const u32 *key)
{
	struct nft_rbtree *priv = nft_set_priv(set);
	unsigned int seq = read_seqcount_begin(&priv->count);
	const struct nft_set_ext *ext;

	ext = __nft_rbtree_lookup(net, set, key, seq);
	if (ext || !read_seqcount_retry(&priv->count, seq))
		return ext;

	read_lock_bh(&priv->lock);
	seq = read_seqcount_begin(&priv->count);
	ext = __nft_rbtree_lookup(net, set, key, seq);
	read_unlock_bh(&priv->lock);

	return ext;
}

static bool __nft_rbtree_get(const struct net *net, const struct nft_set *set,
			     const u32 *key, struct nft_rbtree_elem **elem,
			     unsigned int seq, unsigned int flags, u8 genmask)
{
	struct nft_rbtree_elem *rbe, *interval = NULL;
	struct nft_rbtree *priv = nft_set_priv(set);
	const struct rb_node *parent;
	const void *this;
	int d;

	parent = rcu_dereference_raw(priv->root.rb_node);
	while (parent != NULL) {
		if (read_seqcount_retry(&priv->count, seq))
			return false;

		rbe = rb_entry