linux/net/netfilter/nft_set_rbtree.c
<<
>>
Prefs
   1/*
   2 * Copyright (c) 2008-2009 Patrick McHardy <kaber@trash.net>
   3 *
   4 * This program is free software; you can redistribute it and/or modify
   5 * it under the terms of the GNU General Public License version 2 as
   6 * published by the Free Software Foundation.
   7 *
   8 * Development of this code funded by Astaro AG (http://www.astaro.com/)
   9 */
  10
  11#include <linux/kernel.h>
  12#include <linux/init.h>
  13#include <linux/module.h>
  14#include <linux/list.h>
  15#include <linux/rbtree.h>
  16#include <linux/netlink.h>
  17#include <linux/netfilter.h>
  18#include <linux/netfilter/nf_tables.h>
  19#include <net/netfilter/nf_tables.h>
  20
  21struct nft_rbtree {
  22        rwlock_t                lock;
  23        struct rb_root          root;
  24};
  25
  26struct nft_rbtree_elem {
  27        struct rb_node          node;
  28        struct nft_set_ext      ext;
  29};
  30
  31static bool nft_rbtree_interval_end(const struct nft_rbtree_elem *rbe)
  32{
  33        return nft_set_ext_exists(&rbe->ext, NFT_SET_EXT_FLAGS) &&
  34               (*nft_set_ext_flags(&rbe->ext) & NFT_SET_ELEM_INTERVAL_END);
  35}
  36
  37static bool nft_rbtree_equal(const struct nft_set *set, const void *this,
  38                             const struct nft_rbtree_elem *interval)
  39{
  40        return memcmp(this, nft_set_ext_key(&interval->ext), set->klen) == 0;
  41}
  42
  43static bool nft_rbtree_lookup(const struct net *net, const struct nft_set *set,
  44                              const u32 *key, const struct nft_set_ext **ext)
  45{
  46        struct nft_rbtree *priv = nft_set_priv(set);
  47        const struct nft_rbtree_elem *rbe, *interval = NULL;
  48        u8 genmask = nft_genmask_cur(net);
  49        const struct rb_node *parent;
  50        const void *this;
  51        int d;
  52
  53        read_lock_bh(&priv->lock);
  54        parent = priv->root.rb_node;
  55        while (parent != NULL) {
  56                rbe = rb_entry(parent, struct nft_rbtree_elem, node);
  57
  58                this = nft_set_ext_key(&rbe->ext);
  59                d = memcmp(this, key, set->klen);
  60                if (d < 0) {
  61                        parent = parent->rb_left;
  62                        if (interval &&
  63                            nft_rbtree_equal(set, this, interval) &&
  64                            nft_rbtree_interval_end(this) &&
  65                            !nft_rbtree_interval_end(interval))
  66                                continue;
  67                        interval = rbe;
  68                } else if (d > 0)
  69                        parent = parent->rb_right;
  70                else {
  71                        if (!nft_set_elem_active(&rbe->ext, genmask)) {
  72                                parent = parent->rb_left;
  73                                continue;
  74                        }
  75                        if (nft_rbtree_interval_end(rbe))
  76                                goto out;
  77                        read_unlock_bh(&priv->lock);
  78
  79                        *ext = &rbe->ext;
  80                        return true;
  81                }
  82        }
  83
  84        if (set->flags & NFT_SET_INTERVAL && interval != NULL &&
  85            nft_set_elem_active(&interval->ext, genmask) &&
  86            !nft_rbtree_interval_end(interval)) {
  87                read_unlock_bh(&priv->lock);
  88                *ext = &interval->ext;
  89                return true;
  90        }
  91out:
  92        read_unlock_bh(&priv->lock);
  93        return false;
  94}
  95
  96static int __nft_rbtree_insert(const struct net *net, const struct nft_set *set,
  97                               struct nft_rbtree_elem *new,
  98                               struct nft_set_ext **ext)
  99{
 100        struct nft_rbtree *priv = nft_set_priv(set);
 101        u8 genmask = nft_genmask_next(net);
 102        struct nft_rbtree_elem *rbe;
 103        struct rb_node *parent, **p;
 104        int d;
 105
 106        parent = NULL;
 107        p = &priv->root.rb_node;
 108        while (*p != NULL) {
 109                parent = *p;
 110                rbe = rb_entry(parent, struct nft_rbtree_elem, node);
 111                d = memcmp(nft_set_ext_key(&rbe->ext),
 112                           nft_set_ext_key(&new->ext),
 113                           set->klen);
 114                if (d < 0)
 115                        p = &parent->rb_left;
 116                else if (d > 0)
 117                        p = &parent->rb_right;
 118                else {
 119                        if (nft_rbtree_interval_end(rbe) &&
 120                            !nft_rbtree_interval_end(new)) {
 121                                p = &parent->rb_left;
 122                        } else if (!nft_rbtree_interval_end(rbe) &&
 123                                   nft_rbtree_interval_end(new)) {
 124                                p = &parent->rb_right;
 125                        } else if (nft_set_elem_active(&rbe->ext, genmask)) {
 126                                *ext = &rbe->ext;
 127                                return -EEXIST;
 128                        } else {
 129                                p = &parent->rb_left;
 130                        }
 131                }
 132        }
 133        rb_link_node(&new->node, parent, p);
 134        rb_insert_color(&new->node, &priv->root);
 135        return 0;
 136}
 137
 138static int nft_rbtree_insert(const struct net *net, const struct nft_set *set,
 139                             const struct nft_set_elem *elem,
 140                             struct nft_set_ext **ext)
 141{
 142        struct nft_rbtree *priv = nft_set_priv(set);
 143        struct nft_rbtree_elem *rbe = elem->priv;
 144        int err;
 145
 146        write_lock_bh(&priv->lock);
 147        err = __nft_rbtree_insert(net, set, rbe, ext);
 148        write_unlock_bh(&priv->lock);
 149
 150        return err;
 151}
 152
 153static void nft_rbtree_remove(const struct net *net,
 154                              const struct nft_set *set,
 155                              const struct nft_set_elem *elem)
 156{
 157        struct nft_rbtree *priv = nft_set_priv(set);
 158        struct nft_rbtree_elem *rbe = elem->priv;
 159
 160        write_lock_bh(&priv->lock);
 161        rb_erase(&rbe->node, &priv->root);
 162        write_unlock_bh(&priv->lock);
 163}
 164
 165static void nft_rbtree_activate(const struct net *net,
 166                                const struct nft_set *set,
 167                                const struct nft_set_elem *elem)
 168{
 169        struct nft_rbtree_elem *rbe = elem->priv;
 170
 171        nft_set_elem_change_active(net, set, &rbe->ext);
 172}
 173
 174static bool nft_rbtree_flush(const struct net *net,
 175                             const struct nft_set *set, void *priv)
 176{
 177        struct nft_rbtree_elem *rbe = priv;
 178
 179        nft_set_elem_change_active(net, set, &rbe->ext);
 180        return true;
 181}
 182
 183static void *nft_rbtree_deactivate(const struct net *net,
 184                                   const struct nft_set *set,
 185                                   const struct nft_set_elem *elem)
 186{
 187        const struct nft_rbtree *priv = nft_set_priv(set);
 188        const struct rb_node *parent = priv->root.rb_node;
 189        struct nft_rbtree_elem *rbe, *this = elem->priv;
 190        u8 genmask = nft_genmask_next(net);
 191        int d;
 192
 193        while (parent != NULL) {
 194                rbe = rb_entry(parent, struct nft_rbtree_elem, node);
 195
 196                d = memcmp(nft_set_ext_key(&rbe->ext), &elem->key.val,
 197                                           set->klen);
 198                if (d < 0)
 199                        parent = parent->rb_left;
 200                else if (d > 0)
 201                        parent = parent->rb_right;
 202                else {
 203                        if (!nft_set_elem_active(&rbe->ext, genmask)) {
 204                                parent = parent->rb_left;
 205                                continue;
 206                        }
 207                        if (nft_rbtree_interval_end(rbe) &&
 208                            !nft_rbtree_interval_end(this)) {
 209                                parent = parent->rb_left;
 210                                continue;
 211                        } else if (!nft_rbtree_interval_end(rbe) &&
 212                                   nft_rbtree_interval_end(this)) {
 213                                parent = parent->rb_right;
 214                                continue;
 215                        }
 216                        nft_rbtree_flush(net, set, rbe);
 217                        return rbe;
 218                }
 219        }
 220        return NULL;
 221}
 222
 223static void nft_rbtree_walk(const struct nft_ctx *ctx,
 224                            struct nft_set *set,
 225                            struct nft_set_iter *iter)
 226{
 227        struct nft_rbtree *priv = nft_set_priv(set);
 228        struct nft_rbtree_elem *rbe;
 229        struct nft_set_elem elem;
 230        struct rb_node *node;
 231
 232        read_lock_bh(&priv->lock);
 233        for (node = rb_first(&priv->root); node != NULL; node = rb_next(node)) {
 234                rbe = rb_entry(node, struct nft_rbtree_elem, node);
 235
 236                if (iter->count < iter->skip)
 237                        goto cont;
 238                if (!nft_set_elem_active(&rbe->ext, iter->genmask))
 239                        goto cont;
 240
 241                elem.priv = rbe;
 242
 243                iter->err = iter->fn(ctx, set, iter, &elem);
 244                if (iter->err < 0) {
 245                        read_unlock_bh(&priv->lock);
 246                        return;
 247                }
 248cont:
 249                iter->count++;
 250        }
 251        read_unlock_bh(&priv->lock);
 252}
 253
 254static unsigned int nft_rbtree_privsize(const struct nlattr * const nla[])
 255{
 256        return sizeof(struct nft_rbtree);
 257}
 258
 259static int nft_rbtree_init(const struct nft_set *set,
 260                           const struct nft_set_desc *desc,
 261                           const struct nlattr * const nla[])
 262{
 263        struct nft_rbtree *priv = nft_set_priv(set);
 264
 265        rwlock_init(&priv->lock);
 266        priv->root = RB_ROOT;
 267        return 0;
 268}
 269
 270static void nft_rbtree_destroy(const struct nft_set *set)
 271{
 272        struct nft_rbtree *priv = nft_set_priv(set);
 273        struct nft_rbtree_elem *rbe;
 274        struct rb_node *node;
 275
 276        while ((node = priv->root.rb_node) != NULL) {
 277                rb_erase(node, &priv->root);
 278                rbe = rb_entry(node, struct nft_rbtree_elem, node);
 279                nft_set_elem_destroy(set, rbe, true);
 280        }
 281}
 282
 283static bool nft_rbtree_estimate(const struct nft_set_desc *desc, u32 features,
 284                                struct nft_set_estimate *est)
 285{
 286        unsigned int nsize;
 287
 288        nsize = sizeof(struct nft_rbtree_elem);
 289        if (desc->size)
 290                est->size = sizeof(struct nft_rbtree) + desc->size * nsize;
 291        else
 292                est->size = nsize;
 293
 294        est->lookup = NFT_SET_CLASS_O_LOG_N;
 295        est->space  = NFT_SET_CLASS_O_N;
 296
 297        return true;
 298}
 299
 300static struct nft_set_ops nft_rbtree_ops __read_mostly = {
 301        .privsize       = nft_rbtree_privsize,
 302        .elemsize       = offsetof(struct nft_rbtree_elem, ext),
 303        .estimate       = nft_rbtree_estimate,
 304        .init           = nft_rbtree_init,
 305        .destroy        = nft_rbtree_destroy,
 306        .insert         = nft_rbtree_insert,
 307        .remove         = nft_rbtree_remove,
 308        .deactivate     = nft_rbtree_deactivate,
 309        .flush          = nft_rbtree_flush,
 310        .activate       = nft_rbtree_activate,
 311        .lookup         = nft_rbtree_lookup,
 312        .walk           = nft_rbtree_walk,
 313        .features       = NFT_SET_INTERVAL | NFT_SET_MAP | NFT_SET_OBJECT,
 314        .owner          = THIS_MODULE,
 315};
 316
 317static int __init nft_rbtree_module_init(void)
 318{
 319        return nft_register_set(&nft_rbtree_ops);
 320}
 321
 322static void __exit nft_rbtree_module_exit(void)
 323{
 324        nft_unregister_set(&nft_rbtree_ops);
 325}
 326
 327module_init(nft_rbtree_module_init);
 328module_exit(nft_rbtree_module_exit);
 329
 330MODULE_LICENSE("GPL");
 331MODULE_AUTHOR("Patrick McHardy <kaber@trash.net>");
 332MODULE_ALIAS_NFT_SET();
 333