netfilter: nf_conncount: split count_tree_node rbtree walk into helper

Add find_tree_node() helper that fetches a matching rbtree node.

This is used by followup patch to optionally search the tree again while
preventing concurrent updates via tree lock.

Signed-off-by: Florian Westphal <fw@strlen.de>
Signed-off-by: Pablo Neira Ayuso <pablo@netfilter.org>

diff --git a/net/netfilter/nf_conncount.c b/net/netfilter/nf_conncount.c
index faecc05d34d44f7ce30089e17ebe0a9c2db19a56..56ac64ecfb75d58bd84e345dff99c6bf857d431f 100644 (file)
--- a/net/netfilter/nf_conncount.c
+++ b/net/netfilter/nf_conncount.c
@@ -488,6 +488,34 @@ out_unlock:
        return count;
 }
 
+static struct nf_conncount_rb *
+find_tree_node(struct nf_conncount_root *root, struct nf_conncount_data *data,
+              const u32 *key)
+{
+       struct rb_node *parent;
+
+       parent = rcu_dereference_check(root->root.rb_node,
+                                      lockdep_is_held(&root->lock));
+       while (parent) {
+               struct nf_conncount_rb *rbconn;
+               int diff;
+
+               rbconn = rb_entry(parent, struct nf_conncount_rb, node);
+
+               diff = key_diff(key, rbconn->key, data->keylen);
+               if (diff < 0)
+                       parent = rcu_dereference_check(parent->rb_left,
+                                                      lockdep_is_held(&root->lock));
+               else if (diff > 0)
+                       parent = rcu_dereference_check(parent->rb_right,
+                                                      lockdep_is_held(&root->lock));
+               else
+                       return rbconn;
+       }
+
+       return ERR_PTR(-ENOENT);
+}
+
 static unsigned int
 count_tree(struct net *net,
           const struct sk_buff *skb,
@@ -496,59 +524,53 @@ count_tree(struct net *net,
           const u32 *key)
 {
        struct nf_conncount_root *root;
-       struct rb_node *parent;
        struct nf_conncount_rb *rbconn;
        unsigned int hash;
+       int ret;
 
        hash = jhash2(key, data->keylen, data->initval) % CONNCOUNT_SLOTS;
        root = &data->root[hash];
 
-       parent = rcu_dereference(root->root.rb_node);
-       while (parent) {
-               int diff;
-
-               rbconn = rb_entry(parent, struct nf_conncount_rb, node);
+       rbconn = find_tree_node(root, data, key);
+       if (IS_ERR(rbconn)) {
+               if (PTR_ERR(rbconn) == -ENOENT) {
+                       if (!skb)
+                               return 0;
 
-               diff = key_diff(key, rbconn->key, data->keylen);
-               if (diff < 0) {
-                       parent = rcu_dereference(parent->rb_left);
-               } else if (diff > 0) {
-                       parent = rcu_dereference(parent->rb_right);
-               } else {
-                       int ret;
+                       return insert_tree(net, skb, l3num, data, hash, key);
+               }
+               DEBUG_NET_WARN_ON_ONCE(IS_ERR(rbconn));
+       }
 
-                       if (!skb) {
-                               nf_conncount_gc_list(net, &rbconn->list);
-                               return rbconn->list.count;
-                       }
+       DEBUG_NET_WARN_ON_ONCE(IS_ERR_OR_NULL(rbconn));
+       if (IS_ERR_OR_NULL(rbconn))
+               return 0;
 
-                       spin_lock_bh(&rbconn->list.list_lock);
-                       /* Node might be about to be free'd.
-                        * We need to defer to insert_tree() in this case.
-                        */
-                       if (rbconn->list.count == 0) {
-                               spin_unlock_bh(&rbconn->list.list_lock);
-                               break;
-                       }
+       if (!skb) {
+               nf_conncount_gc_list(net, &rbconn->list);
+               return rbconn->list.count;
+       }
 
-                       /* same source network -> be counted! */
-                       ret = __nf_conncount_add(net, skb, l3num, &rbconn->list);
-                       spin_unlock_bh(&rbconn->list.list_lock);
-                       if (ret && ret != -EEXIST) {
-                               return 0; /* hotdrop */
-                       } else {
-                               /* -EEXIST means add was skipped, update the list */
-                               if (ret == -EEXIST)
-                                       nf_conncount_gc_list(net, &rbconn->list);
-                               return rbconn->list.count;
-                       }
-               }
+       spin_lock_bh(&rbconn->list.list_lock);
+       /* Node might be about to be free'd.
+        * We need to defer to insert_tree() in this case.
+        */
+       if (rbconn->list.count == 0) {
+               spin_unlock_bh(&rbconn->list.list_lock);
+               return insert_tree(net, skb, l3num, data, hash, key);
        }
 
-       if (!skb)
-               return 0;
+       /* same source network -> be counted! */
+       ret = __nf_conncount_add(net, skb, l3num, &rbconn->list);
+       spin_unlock_bh(&rbconn->list.list_lock);
+
+       if (ret && ret != -EEXIST)
+               return 0; /* hotdrop */
+       /* -EEXIST means add was skipped, update the list */
+       if (ret == -EEXIST)
+               nf_conncount_gc_list(net, &rbconn->list);
 
-       return insert_tree(net, skb, l3num, data, hash, key);
+       return rbconn->list.count;
 }
 
 static void tree_gc_worker(struct work_struct *work)