net/netfilter/nft_rbtree.c

   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
  21 static DEFINE_SPINLOCK(nft_rbtree_lock);
  22
  23 struct nft_rbtree {
  24         struct rb_root          root;
  25 };
  26
  27 struct nft_rbtree_elem {
  28         struct rb_node          node;
  29         struct nft_set_ext      ext;
  30 };
  31
  32
  33 static bool nft_rbtree_lookup(const struct nft_set *set, const u32 *key,
  34                               const struct nft_set_ext **ext)
  35 {
  36         const struct nft_rbtree *priv = nft_set_priv(set);
  37         const struct nft_rbtree_elem *rbe, *interval = NULL;
  38         const struct rb_node *parent;
  39         u8 genmask = nft_genmask_cur(read_pnet(&set->pnet));
  40         int d;
  41
  42         spin_lock_bh(&nft_rbtree_lock);
  43         parent = priv->root.rb_node;
  44         while (parent != NULL) {
  45                 rbe = rb_entry(parent, struct nft_rbtree_elem, node);
  46
  47                 d = memcmp(nft_set_ext_key(&rbe->ext), key, set->klen);
  48                 if (d < 0) {
  49                         parent = parent->rb_left;
  50                         interval = rbe;
  51                 } else if (d > 0)
  52                         parent = parent->rb_right;
  53                 else {
  54 found:
  55                         if (!nft_set_elem_active(&rbe->ext, genmask)) {
  56                                 parent = parent->rb_left;
  57                                 continue;
  58                         }
  59                         if (nft_set_ext_exists(&rbe->ext, NFT_SET_EXT_FLAGS) &&
  60                             *nft_set_ext_flags(&rbe->ext) &
  61                             NFT_SET_ELEM_INTERVAL_END)
  62                                 goto out;
  63                         spin_unlock_bh(&nft_rbtree_lock);
  64
  65                         *ext = &rbe->ext;
  66                         return true;
  67                 }
  68         }
  69
  70         if (set->flags & NFT_SET_INTERVAL && interval != NULL) {
  71                 rbe = interval;
  72                 goto found;
  73         }
  74 out:
  75         spin_unlock_bh(&nft_rbtree_lock);
  76         return false;
  77 }
  78
  79 static int __nft_rbtree_insert(const struct nft_set *set,
  80                                struct nft_rbtree_elem *new)
  81 {
  82         struct nft_rbtree *priv = nft_set_priv(set);
  83         struct nft_rbtree_elem *rbe;
  84         struct rb_node *parent, **p;
  85         u8 genmask = nft_genmask_next(read_pnet(&set->pnet));
  86         int d;
  87
  88         parent = NULL;
  89         p = &priv->root.rb_node;
  90         while (*p != NULL) {
  91                 parent = *p;
  92                 rbe = rb_entry(parent, struct nft_rbtree_elem, node);
  93                 d = memcmp(nft_set_ext_key(&rbe->ext),
  94                            nft_set_ext_key(&new->ext),
  95                            set->klen);
  96                 if (d < 0)
  97                         p = &parent->rb_left;
  98                 else if (d > 0)
  99                         p = &parent->rb_right;
 100                 else {
 101                         if (nft_set_elem_active(&rbe->ext, genmask))
 102                                 return -EEXIST;
 103                         p = &parent->rb_left;
 104                 }
 105         }
 106         rb_link_node(&new->node, parent, p);
 107         rb_insert_color(&new->node, &priv->root);
 108         return 0;
 109 }
 110
 111 static int nft_rbtree_insert(const struct nft_set *set,
 112                              const struct nft_set_elem *elem)
 113 {
 114         struct nft_rbtree_elem *rbe = elem->priv;
 115         int err;
 116
 117         spin_lock_bh(&nft_rbtree_lock);
 118         err = __nft_rbtree_insert(set, rbe);
 119         spin_unlock_bh(&nft_rbtree_lock);
 120
 121         return err;
 122 }
 123
 124 static void nft_rbtree_remove(const struct nft_set *set,
 125                               const struct nft_set_elem *elem)
 126 {
 127         struct nft_rbtree *priv = nft_set_priv(set);
 128         struct nft_rbtree_elem *rbe = elem->priv;
 129
 130         spin_lock_bh(&nft_rbtree_lock);
 131         rb_erase(&rbe->node, &priv->root);
 132         spin_unlock_bh(&nft_rbtree_lock);
 133 }
 134
 135 static void nft_rbtree_activate(const struct nft_set *set,
 136                                 const struct nft_set_elem *elem)
 137 {
 138         struct nft_rbtree_elem *rbe = elem->priv;
 139
 140         nft_set_elem_change_active(set, &rbe->ext);
 141 }
 142
 143 static void *nft_rbtree_deactivate(const struct nft_set *set,
 144                                    const struct nft_set_elem *elem)
 145 {
 146         const struct nft_rbtree *priv = nft_set_priv(set);
 147         const struct rb_node *parent = priv->root.rb_node;
 148         struct nft_rbtree_elem *rbe;
 149         u8 genmask = nft_genmask_cur(read_pnet(&set->pnet));
 150         int d;
 151
 152         while (parent != NULL) {
 153                 rbe = rb_entry(parent, struct nft_rbtree_elem, node);
 154
 155                 d = memcmp(nft_set_ext_key(&rbe->ext), &elem->key.val,
 156                                            set->klen);
 157                 if (d < 0)
 158                         parent = parent->rb_left;
 159                 else if (d > 0)
 160                         parent = parent->rb_right;
 161                 else {
 162                         if (!nft_set_elem_active(&rbe->ext, genmask)) {
 163                                 parent = parent->rb_left;
 164                                 continue;
 165                         }
 166                         nft_set_elem_change_active(set, &rbe->ext);
 167                         return rbe;
 168                 }
 169         }
 170         return NULL;
 171 }
 172
 173 static void nft_rbtree_walk(const struct nft_ctx *ctx,
 174                             const struct nft_set *set,
 175                             struct nft_set_iter *iter)
 176 {
 177         const struct nft_rbtree *priv = nft_set_priv(set);
 178         struct nft_rbtree_elem *rbe;
 179         struct nft_set_elem elem;
 180         struct rb_node *node;
 181         u8 genmask = nft_genmask_cur(read_pnet(&set->pnet));
 182
 183         spin_lock_bh(&nft_rbtree_lock);
 184         for (node = rb_first(&priv->root); node != NULL; node = rb_next(node)) {
 185                 rbe = rb_entry(node, struct nft_rbtree_elem, node);
 186
 187                 if (iter->count < iter->skip)
 188                         goto cont;
 189                 if (!nft_set_elem_active(&rbe->ext, genmask))
 190                         goto cont;
 191
 192                 elem.priv = rbe;
 193
 194                 iter->err = iter->fn(ctx, set, iter, &elem);
 195                 if (iter->err < 0) {
 196                         spin_unlock_bh(&nft_rbtree_lock);
 197                         return;
 198                 }
 199 cont:
 200                 iter->count++;
 201         }
 202         spin_unlock_bh(&nft_rbtree_lock);
 203 }
 204
 205 static unsigned int nft_rbtree_privsize(const struct nlattr * const nla[])
 206 {
 207         return sizeof(struct nft_rbtree);
 208 }
 209
 210 static int nft_rbtree_init(const struct nft_set *set,
 211                            const struct nft_set_desc *desc,
 212                            const struct nlattr * const nla[])
 213 {
 214         struct nft_rbtree *priv = nft_set_priv(set);
 215
 216         priv->root = RB_ROOT;
 217         return 0;
 218 }
 219
 220 static void nft_rbtree_destroy(const struct nft_set *set)
 221 {
 222         struct nft_rbtree *priv = nft_set_priv(set);
 223         struct nft_rbtree_elem *rbe;
 224         struct rb_node *node;
 225
 226         while ((node = priv->root.rb_node) != NULL) {
 227                 rb_erase(node, &priv->root);
 228                 rbe = rb_entry(node, struct nft_rbtree_elem, node);
 229                 nft_set_elem_destroy(set, rbe);
 230         }
 231 }
 232
 233 static bool nft_rbtree_estimate(const struct nft_set_desc *desc, u32 features,
 234                                 struct nft_set_estimate *est)
 235 {
 236         unsigned int nsize;
 237
 238         nsize = sizeof(struct nft_rbtree_elem);
 239         if (desc->size)
 240                 est->size = sizeof(struct nft_rbtree) + desc->size * nsize;
 241         else
 242                 est->size = nsize;
 243
 244         est->class = NFT_SET_CLASS_O_LOG_N;
 245
 246         return true;
 247 }
 248
 249 static struct nft_set_ops nft_rbtree_ops __read_mostly = {
 250         .privsize       = nft_rbtree_privsize,
 251         .elemsize       = offsetof(struct nft_rbtree_elem, ext),
 252         .estimate       = nft_rbtree_estimate,
 253         .init           = nft_rbtree_init,
 254         .destroy        = nft_rbtree_destroy,
 255         .insert         = nft_rbtree_insert,
 256         .remove         = nft_rbtree_remove,
 257         .deactivate     = nft_rbtree_deactivate,
 258         .activate       = nft_rbtree_activate,
 259         .lookup         = nft_rbtree_lookup,
 260         .walk           = nft_rbtree_walk,
 261         .features       = NFT_SET_INTERVAL | NFT_SET_MAP,
 262         .owner          = THIS_MODULE,
 263 };
 264
 265 static int __init nft_rbtree_module_init(void)
 266 {
 267         return nft_register_set(&nft_rbtree_ops);
 268 }
 269
 270 static void __exit nft_rbtree_module_exit(void)
 271 {
 272         nft_unregister_set(&nft_rbtree_ops);
 273 }
 274
 275 module_init(nft_rbtree_module_init);
 276 module_exit(nft_rbtree_module_exit);
 277
 278 MODULE_LICENSE("GPL");
 279 MODULE_AUTHOR("Patrick McHardy <kaber@trash.net>");
 280 MODULE_ALIAS_NFT_SET();