xref: /linux/net/netfilter/nf_conncount.c (revision 49bda4826843be0ef97a162009a29ea3a63f3935)
1 // SPDX-License-Identifier: GPL-2.0-only
2 /*
3  * count the number of connections matching an arbitrary key.
4  *
5  * (C) 2017 Red Hat GmbH
6  * Author: Florian Westphal <fw@strlen.de>
7  *
8  * split from xt_connlimit.c:
9  *   (c) 2000 Gerd Knorr <kraxel@bytesex.org>
10  *   Nov 2002: Martin Bene <martin.bene@icomedias.com>:
11  *		only ignore TIME_WAIT or gone connections
12  *   (C) CC Computer Consultants GmbH, 2007
13  */
14 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
15 #include <linux/in.h>
16 #include <linux/in6.h>
17 #include <linux/ip.h>
18 #include <linux/ipv6.h>
19 #include <linux/jhash.h>
20 #include <linux/slab.h>
21 #include <linux/list.h>
22 #include <linux/rbtree.h>
23 #include <linux/module.h>
24 #include <linux/random.h>
25 #include <linux/skbuff.h>
26 #include <linux/spinlock.h>
27 #include <linux/netfilter/nf_conntrack_tcp.h>
28 #include <linux/netfilter/x_tables.h>
29 #include <net/netfilter/nf_conntrack.h>
30 #include <net/netfilter/nf_conntrack_count.h>
31 #include <net/netfilter/nf_conntrack_core.h>
32 #include <net/netfilter/nf_conntrack_tuple.h>
33 #include <net/netfilter/nf_conntrack_zones.h>
34 
35 #define CONNCOUNT_SLOTS		256U
36 
37 #define CONNCOUNT_GC_MAX_NODES		8
38 #define CONNCOUNT_GC_MAX_COLLECT	64
39 #define MAX_KEYLEN			5
40 
41 /* we will save the tuples of all connections we care about */
42 struct nf_conncount_tuple {
43 	struct list_head		node;
44 	struct nf_conntrack_tuple	tuple;
45 	struct nf_conntrack_zone	zone;
46 	int				cpu;
47 	u32				jiffies32;
48 };
49 
50 struct nf_conncount_rb {
51 	struct rb_node node;
52 	struct nf_conncount_list list;
53 	u32 key[MAX_KEYLEN];
54 	struct rcu_head rcu_head;
55 };
56 
57 struct nf_conncount_root {
58 	struct rb_root root;
59 	spinlock_t lock;
60 	seqcount_spinlock_t count;
61 };
62 
63 struct nf_conncount_data {
64 	unsigned int keylen;
65 	u32 initval;
66 	struct nf_conncount_root root[CONNCOUNT_SLOTS];
67 	struct net *net;
68 	struct work_struct gc_work;
69 	unsigned long pending_trees[BITS_TO_LONGS(CONNCOUNT_SLOTS)];
70 	unsigned int gc_tree;
71 };
72 
73 static struct kmem_cache *conncount_rb_cachep __read_mostly;
74 static struct kmem_cache *conncount_conn_cachep __read_mostly;
75 
already_closed(const struct nf_conn * conn)76 static inline bool already_closed(const struct nf_conn *conn)
77 {
78 	if (nf_ct_protonum(conn) == IPPROTO_TCP)
79 		return conn->proto.tcp.state == TCP_CONNTRACK_TIME_WAIT ||
80 		       conn->proto.tcp.state == TCP_CONNTRACK_CLOSE;
81 	else
82 		return false;
83 }
84 
key_diff(const u32 * a,const u32 * b,unsigned int klen)85 static int key_diff(const u32 *a, const u32 *b, unsigned int klen)
86 {
87 	return memcmp(a, b, klen * sizeof(u32));
88 }
89 
conn_free(struct nf_conncount_list * list,struct nf_conncount_tuple * conn)90 static void conn_free(struct nf_conncount_list *list,
91 		      struct nf_conncount_tuple *conn)
92 {
93 	lockdep_assert_held(&list->list_lock);
94 
95 	list->count--;
96 	list_del(&conn->node);
97 
98 	kmem_cache_free(conncount_conn_cachep, conn);
99 }
100 
101 static const struct nf_conntrack_tuple_hash *
find_or_evict(struct net * net,struct nf_conncount_list * list,struct nf_conncount_tuple * conn)102 find_or_evict(struct net *net, struct nf_conncount_list *list,
103 	      struct nf_conncount_tuple *conn)
104 {
105 	const struct nf_conntrack_tuple_hash *found;
106 	unsigned long a, b;
107 	int cpu = raw_smp_processor_id();
108 	u32 age;
109 
110 	found = nf_conntrack_find_get(net, &conn->zone, &conn->tuple);
111 	if (found)
112 		return found;
113 	b = conn->jiffies32;
114 	a = (u32)jiffies;
115 
116 	/* conn might have been added just before by another cpu and
117 	 * might still be unconfirmed.  In this case, nf_conntrack_find()
118 	 * returns no result.  Thus only evict if this cpu added the
119 	 * stale entry or if the entry is older than two jiffies.
120 	 */
121 	age = a - b;
122 	if (conn->cpu == cpu || age >= 2) {
123 		conn_free(list, conn);
124 		return ERR_PTR(-ENOENT);
125 	}
126 
127 	return ERR_PTR(-EAGAIN);
128 }
129 
get_ct_or_tuple_from_skb(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conn ** ct,struct nf_conntrack_tuple * tuple,const struct nf_conntrack_zone ** zone,bool * refcounted)130 static bool get_ct_or_tuple_from_skb(struct net *net,
131 				     const struct sk_buff *skb,
132 				     u16 l3num,
133 				     struct nf_conn **ct,
134 				     struct nf_conntrack_tuple *tuple,
135 				     const struct nf_conntrack_zone **zone,
136 				     bool *refcounted)
137 {
138 	const struct nf_conntrack_tuple_hash *h;
139 	enum ip_conntrack_info ctinfo;
140 	struct nf_conn *found_ct;
141 
142 	found_ct = nf_ct_get(skb, &ctinfo);
143 	if (found_ct && !nf_ct_is_template(found_ct)) {
144 		*tuple = found_ct->tuplehash[IP_CT_DIR_ORIGINAL].tuple;
145 		*zone = nf_ct_zone(found_ct);
146 		*ct = found_ct;
147 		return true;
148 	}
149 
150 	if (!nf_ct_get_tuplepr(skb, skb_network_offset(skb), l3num, net, tuple))
151 		return false;
152 
153 	if (found_ct)
154 		*zone = nf_ct_zone(found_ct);
155 
156 	h = nf_conntrack_find_get(net, *zone, tuple);
157 	if (!h)
158 		return true;
159 
160 	found_ct = nf_ct_tuplehash_to_ctrack(h);
161 	*tuple = found_ct->tuplehash[IP_CT_DIR_ORIGINAL].tuple;
162 	*zone = nf_ct_zone(found_ct);
163 	*refcounted = true;
164 	*ct = found_ct;
165 
166 	return true;
167 }
168 
__nf_conncount_add(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conncount_list * list)169 static int __nf_conncount_add(struct net *net,
170 			      const struct sk_buff *skb,
171 			      u16 l3num,
172 			      struct nf_conncount_list *list)
173 {
174 	const struct nf_conntrack_zone *zone = &nf_ct_zone_dflt;
175 	const struct nf_conntrack_tuple_hash *found;
176 	struct nf_conncount_tuple *conn, *conn_n;
177 	struct nf_conntrack_tuple tuple;
178 	struct nf_conn *ct = NULL;
179 	struct nf_conn *found_ct;
180 	unsigned int collect = 0;
181 	bool refcounted = false;
182 	int err = 0;
183 
184 	if (!get_ct_or_tuple_from_skb(net, skb, l3num, &ct, &tuple, &zone, &refcounted))
185 		return -ENOENT;
186 
187 	if (ct && nf_ct_is_confirmed(ct)) {
188 		/* Connection is confirmed but might still be in the setup phase.
189 		 * Only skip the tracking if it is fully assured. This guarantees
190 		 * that setup packets or retransmissions are properly counted and
191 		 * deduplicated.
192 		 */
193 		if (test_bit(IPS_ASSURED_BIT, &ct->status)) {
194 			err = -EEXIST;
195 			goto out_put;
196 		}
197 
198 		goto check_connections;
199 	}
200 
201 	if ((u32)jiffies == list->last_gc &&
202 	    (list->count - list->last_gc_count) < CONNCOUNT_GC_MAX_COLLECT)
203 		goto add_new_node;
204 
205 check_connections:
206 	/* check the saved connections */
207 	list_for_each_entry_safe(conn, conn_n, &list->head, node) {
208 		if (collect > CONNCOUNT_GC_MAX_COLLECT)
209 			break;
210 
211 		found = find_or_evict(net, list, conn);
212 		if (IS_ERR(found)) {
213 			/* Not found, but might be about to be confirmed */
214 			if (PTR_ERR(found) == -EAGAIN) {
215 				if (nf_ct_tuple_equal(&conn->tuple, &tuple) &&
216 				    nf_ct_zone_id(&conn->zone, IP_CT_DIR_ORIGINAL) ==
217 				    nf_ct_zone_id(zone, IP_CT_DIR_ORIGINAL))
218 					goto out_put; /* already exists */
219 			} else {
220 				collect++;
221 			}
222 			continue;
223 		}
224 
225 		found_ct = nf_ct_tuplehash_to_ctrack(found);
226 
227 		if (nf_ct_tuple_equal(&conn->tuple, &tuple) &&
228 		    nf_ct_zone_equal(found_ct, zone, IP_CT_DIR_ORIGINAL)) {
229 			/*
230 			 * We should not see tuples twice unless someone hooks
231 			 * this into a table without "-p tcp --syn".
232 			 *
233 			 * Attempt to avoid a re-add in this case.
234 			 */
235 			nf_ct_put(found_ct);
236 			goto out_put;
237 		} else if (already_closed(found_ct)) {
238 			/*
239 			 * we do not care about connections which are
240 			 * closed already -> ditch it
241 			 */
242 			nf_ct_put(found_ct);
243 			conn_free(list, conn);
244 			collect++;
245 			continue;
246 		}
247 
248 		nf_ct_put(found_ct);
249 	}
250 	list->last_gc = (u32)jiffies;
251 	list->last_gc_count = list->count;
252 
253 add_new_node:
254 	if (unlikely(list->count > INT_MAX)) {
255 		DEBUG_NET_WARN_ON_ONCE(1);
256 		err = -EOVERFLOW;
257 		goto out_put;
258 	}
259 
260 	conn = kmem_cache_alloc(conncount_conn_cachep, GFP_ATOMIC);
261 	if (conn == NULL) {
262 		err = -ENOMEM;
263 		goto out_put;
264 	}
265 
266 	conn->tuple = tuple;
267 	conn->zone = *zone;
268 	conn->cpu = raw_smp_processor_id();
269 	conn->jiffies32 = (u32)jiffies;
270 	list_add_tail(&conn->node, &list->head);
271 	list->count++;
272 
273 out_put:
274 	if (refcounted)
275 		nf_ct_put(ct);
276 	return err;
277 }
278 
nf_conncount_add_skb(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conncount_list * list)279 int nf_conncount_add_skb(struct net *net,
280 			 const struct sk_buff *skb,
281 			 u16 l3num,
282 			 struct nf_conncount_list *list)
283 {
284 	int ret;
285 
286 	/* check the saved connections */
287 	spin_lock_bh(&list->list_lock);
288 	ret = __nf_conncount_add(net, skb, l3num, list);
289 	spin_unlock_bh(&list->list_lock);
290 
291 	return ret;
292 }
293 EXPORT_SYMBOL_GPL(nf_conncount_add_skb);
294 
nf_conncount_list_init(struct nf_conncount_list * list)295 void nf_conncount_list_init(struct nf_conncount_list *list)
296 {
297 	spin_lock_init(&list->list_lock);
298 	INIT_LIST_HEAD(&list->head);
299 	list->count = 0;
300 	list->last_gc_count = 0;
301 	list->last_gc = (u32)jiffies;
302 }
303 EXPORT_SYMBOL_GPL(nf_conncount_list_init);
304 
305 /* Return true if the list is empty. Must be called with BH disabled. */
__nf_conncount_gc_list(struct net * net,struct nf_conncount_list * list)306 static bool __nf_conncount_gc_list(struct net *net,
307 				   struct nf_conncount_list *list)
308 {
309 	const struct nf_conntrack_tuple_hash *found;
310 	struct nf_conncount_tuple *conn, *conn_n;
311 	struct nf_conn *found_ct;
312 	unsigned int collected = 0;
313 	bool ret = false;
314 
315 	/* don't bother if we just did GC */
316 	if ((u32)jiffies == READ_ONCE(list->last_gc))
317 		return false;
318 
319 	list_for_each_entry_safe(conn, conn_n, &list->head, node) {
320 		found = find_or_evict(net, list, conn);
321 		if (IS_ERR(found)) {
322 			if (PTR_ERR(found) == -ENOENT)
323 				collected++;
324 			continue;
325 		}
326 
327 		found_ct = nf_ct_tuplehash_to_ctrack(found);
328 		if (already_closed(found_ct)) {
329 			/*
330 			 * we do not care about connections which are
331 			 * closed already -> ditch it
332 			 */
333 			nf_ct_put(found_ct);
334 			conn_free(list, conn);
335 			collected++;
336 			continue;
337 		}
338 
339 		nf_ct_put(found_ct);
340 		if (collected > CONNCOUNT_GC_MAX_COLLECT)
341 			break;
342 	}
343 
344 	if (!list->count)
345 		ret = true;
346 	list->last_gc = (u32)jiffies;
347 	list->last_gc_count = list->count;
348 
349 	return ret;
350 }
351 
nf_conncount_gc_list(struct net * net,struct nf_conncount_list * list)352 bool nf_conncount_gc_list(struct net *net,
353 			  struct nf_conncount_list *list)
354 {
355 	bool ret;
356 
357 	/* don't bother if other cpu is already doing GC */
358 	if (!spin_trylock_bh(&list->list_lock))
359 		return false;
360 
361 	ret = __nf_conncount_gc_list(net, list);
362 	spin_unlock_bh(&list->list_lock);
363 
364 	return ret;
365 }
366 EXPORT_SYMBOL_GPL(nf_conncount_gc_list);
367 
__tree_nodes_free(struct rcu_head * h)368 static void __tree_nodes_free(struct rcu_head *h)
369 {
370 	struct nf_conncount_rb *rbconn;
371 
372 	rbconn = container_of(h, struct nf_conncount_rb, rcu_head);
373 	kmem_cache_free(conncount_rb_cachep, rbconn);
374 }
375 
tree_nodes_free(struct nf_conncount_root * root,struct nf_conncount_rb * gc_nodes[],unsigned int gc_count)376 static void tree_nodes_free(struct nf_conncount_root *root,
377 			    struct nf_conncount_rb *gc_nodes[],
378 			    unsigned int gc_count)
379 {
380 	struct nf_conncount_rb *rbconn;
381 
382 	lockdep_assert_held(&root->lock);
383 
384 	while (gc_count) {
385 		rbconn = gc_nodes[--gc_count];
386 		spin_lock(&rbconn->list.list_lock);
387 		if (!rbconn->list.count) {
388 			write_seqcount_begin(&root->count);
389 			rb_erase(&rbconn->node, &root->root);
390 			call_rcu(&rbconn->rcu_head, __tree_nodes_free);
391 			write_seqcount_end(&root->count);
392 		}
393 		spin_unlock(&rbconn->list.list_lock);
394 	}
395 }
396 
schedule_gc_worker(struct nf_conncount_data * data,int tree)397 static void schedule_gc_worker(struct nf_conncount_data *data, int tree)
398 {
399 	set_bit(tree, data->pending_trees);
400 	schedule_work(&data->gc_work);
401 }
402 
403 static unsigned int
insert_tree(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conncount_data * data,unsigned int hash,const u32 * key)404 insert_tree(struct net *net,
405 	    const struct sk_buff *skb,
406 	    u16 l3num,
407 	    struct nf_conncount_data *data,
408 	    unsigned int hash,
409 	    const u32 *key)
410 {
411 	struct nf_conncount_root *root = &data->root[hash];
412 	struct nf_conncount_rb *gc_nodes[CONNCOUNT_GC_MAX_NODES];
413 	const struct nf_conntrack_zone *zone = &nf_ct_zone_dflt;
414 	bool do_gc = true, refcounted = false;
415 	unsigned int count = 0, gc_count = 0;
416 	struct rb_node **rbnode, *parent;
417 	struct nf_conntrack_tuple tuple;
418 	struct nf_conncount_tuple *conn;
419 	struct nf_conncount_rb *rbconn;
420 	struct nf_conn *ct = NULL;
421 
422 	spin_lock_bh(&root->lock);
423 restart:
424 	parent = NULL;
425 	rbnode = &root->root.rb_node;
426 	while (*rbnode) {
427 		int diff;
428 		rbconn = rb_entry(*rbnode, struct nf_conncount_rb, node);
429 
430 		parent = *rbnode;
431 		diff = key_diff(key, rbconn->key, data->keylen);
432 		if (diff < 0) {
433 			rbnode = &((*rbnode)->rb_left);
434 		} else if (diff > 0) {
435 			rbnode = &((*rbnode)->rb_right);
436 		} else {
437 			int ret;
438 
439 			ret = nf_conncount_add_skb(net, skb, l3num, &rbconn->list);
440 			if (ret && ret != -EEXIST)
441 				count = 0; /* hotdrop */
442 			else
443 				count = rbconn->list.count;
444 			tree_nodes_free(root, gc_nodes, gc_count);
445 			goto out_unlock;
446 		}
447 
448 		if (gc_count >= ARRAY_SIZE(gc_nodes))
449 			continue;
450 
451 		if (do_gc && nf_conncount_gc_list(net, &rbconn->list))
452 			gc_nodes[gc_count++] = rbconn;
453 	}
454 
455 	if (gc_count) {
456 		tree_nodes_free(root, gc_nodes, gc_count);
457 		schedule_gc_worker(data, hash);
458 		gc_count = 0;
459 		do_gc = false;
460 		goto restart;
461 	}
462 
463 	if (get_ct_or_tuple_from_skb(net, skb, l3num, &ct, &tuple, &zone, &refcounted)) {
464 		/* expected case: match, insert new node */
465 		rbconn = kmem_cache_alloc(conncount_rb_cachep, GFP_ATOMIC);
466 		if (rbconn == NULL)
467 			goto out_unlock;
468 
469 		conn = kmem_cache_alloc(conncount_conn_cachep, GFP_ATOMIC);
470 		if (conn == NULL) {
471 			kmem_cache_free(conncount_rb_cachep, rbconn);
472 			goto out_unlock;
473 		}
474 
475 		conn->tuple = tuple;
476 		conn->zone = *zone;
477 		conn->cpu = raw_smp_processor_id();
478 		conn->jiffies32 = (u32)jiffies;
479 		memcpy(rbconn->key, key, sizeof(u32) * data->keylen);
480 
481 		nf_conncount_list_init(&rbconn->list);
482 		list_add(&conn->node, &rbconn->list.head);
483 		count = 1;
484 		rbconn->list.count = count;
485 
486 		write_seqcount_begin(&root->count);
487 		rb_link_node_rcu(&rbconn->node, parent, rbnode);
488 		rb_insert_color(&rbconn->node, &root->root);
489 		write_seqcount_end(&root->count);
490 	}
491 out_unlock:
492 	if (refcounted)
493 		nf_ct_put(ct);
494 	spin_unlock_bh(&root->lock);
495 	return count;
496 }
497 
498 static struct nf_conncount_rb *
find_tree_node(struct nf_conncount_root * root,struct nf_conncount_data * data,const u32 * key)499 find_tree_node(struct nf_conncount_root *root, struct nf_conncount_data *data,
500 	       const u32 *key)
501 {
502 	unsigned int seq = read_seqcount_begin(&root->count);
503 	struct rb_node *parent;
504 
505 	parent = rcu_dereference_check(root->root.rb_node,
506 				       lockdep_is_held(&root->lock));
507 	while (parent) {
508 		struct nf_conncount_rb *rbconn;
509 		int diff;
510 
511 		rbconn = rb_entry(parent, struct nf_conncount_rb, node);
512 
513 		diff = key_diff(key, rbconn->key, data->keylen);
514 		if (diff < 0)
515 			parent = rcu_dereference_check(parent->rb_left,
516 						       lockdep_is_held(&root->lock));
517 		else if (diff > 0)
518 			parent = rcu_dereference_check(parent->rb_right,
519 						       lockdep_is_held(&root->lock));
520 		else
521 			return rbconn;
522 
523 		if (read_seqcount_retry(&root->count, seq))
524 			return ERR_PTR(-EAGAIN);
525 	}
526 
527 	if (read_seqcount_retry(&root->count, seq))
528 		return ERR_PTR(-EAGAIN);
529 
530 	return ERR_PTR(-ENOENT);
531 }
532 
533 static unsigned int
count_tree(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conncount_data * data,const u32 * key)534 count_tree(struct net *net,
535 	   const struct sk_buff *skb,
536 	   u16 l3num,
537 	   struct nf_conncount_data *data,
538 	   const u32 *key)
539 {
540 	struct nf_conncount_root *root;
541 	struct nf_conncount_rb *rbconn;
542 	unsigned int hash;
543 	int ret;
544 
545 	hash = jhash2(key, data->keylen, data->initval) % CONNCOUNT_SLOTS;
546 	root = &data->root[hash];
547 
548 	rbconn = find_tree_node(root, data, key);
549 	if (IS_ERR(rbconn)) {
550 		if (PTR_ERR(rbconn) == -EAGAIN) {
551 			spin_lock_bh(&root->lock);
552 			rbconn = find_tree_node(root, data, key);
553 			spin_unlock_bh(&root->lock);
554 		}
555 
556 		if (PTR_ERR(rbconn) == -ENOENT) {
557 			if (!skb)
558 				return 0;
559 
560 			return insert_tree(net, skb, l3num, data, hash, key);
561 		}
562 		DEBUG_NET_WARN_ON_ONCE(IS_ERR(rbconn));
563 	}
564 
565 	DEBUG_NET_WARN_ON_ONCE(IS_ERR_OR_NULL(rbconn));
566 	if (IS_ERR_OR_NULL(rbconn))
567 		return 0;
568 
569 	if (!skb) {
570 		nf_conncount_gc_list(net, &rbconn->list);
571 		return rbconn->list.count;
572 	}
573 
574 	spin_lock_bh(&rbconn->list.list_lock);
575 	/* Node might be about to be free'd.
576 	 * We need to defer to insert_tree() in this case.
577 	 */
578 	if (rbconn->list.count == 0) {
579 		spin_unlock_bh(&rbconn->list.list_lock);
580 		return insert_tree(net, skb, l3num, data, hash, key);
581 	}
582 
583 	/* same source network -> be counted! */
584 	ret = __nf_conncount_add(net, skb, l3num, &rbconn->list);
585 	spin_unlock_bh(&rbconn->list.list_lock);
586 
587 	if (ret && ret != -EEXIST)
588 		return 0; /* hotdrop */
589 	/* -EEXIST means add was skipped, update the list */
590 	if (ret == -EEXIST)
591 		nf_conncount_gc_list(net, &rbconn->list);
592 
593 	return rbconn->list.count;
594 }
595 
tree_gc_worker(struct work_struct * work)596 static void tree_gc_worker(struct work_struct *work)
597 {
598 	struct nf_conncount_data *data = container_of(work, struct nf_conncount_data, gc_work);
599 	struct nf_conncount_rb *gc_nodes[CONNCOUNT_GC_MAX_NODES], *rbconn;
600 	unsigned int tree, next_tree, gc_count = 0;
601 	struct nf_conncount_root *root;
602 	struct rb_node *node;
603 
604 	if (data->gc_tree == 0)
605 		data->gc_tree = find_first_bit(data->pending_trees, CONNCOUNT_SLOTS);
606 
607 	tree = data->gc_tree % CONNCOUNT_SLOTS;
608 	root = &data->root[tree];
609 
610 	spin_lock_bh(&root->lock);
611 	gc_count = 0;
612 	node = rb_first(&root->root);
613 	while (node != NULL) {
614 		u32 key[MAX_KEYLEN];
615 		bool drop_lock;
616 
617 		rbconn = rb_entry(node, struct nf_conncount_rb, node);
618 		node = rb_next(node);
619 
620 		if (nf_conncount_gc_list(data->net, &rbconn->list))
621 			gc_nodes[gc_count++] = rbconn;
622 
623 		drop_lock = need_resched();
624 
625 		if (drop_lock || gc_count >= ARRAY_SIZE(gc_nodes)) {
626 			tree_nodes_free(root, gc_nodes, gc_count);
627 			gc_count = 0;
628 		}
629 
630 		if (!drop_lock || !node)
631 			continue;
632 
633 		rbconn = rb_entry(node, struct nf_conncount_rb, node);
634 		memcpy(key, rbconn->key, sizeof(key));
635 		spin_unlock_bh(&root->lock);
636 
637 		cond_resched();
638 
639 		spin_lock_bh(&root->lock);
640 		rbconn = find_tree_node(root, data, key);
641 		if (IS_ERR_OR_NULL(rbconn)) /* rbconn was reaped */
642 			break;
643 
644 		node = &rbconn->node;
645 	}
646 
647 	tree_nodes_free(root, gc_nodes, gc_count);
648 	clear_bit(tree, data->pending_trees);
649 
650 	next_tree = (tree + 1) % CONNCOUNT_SLOTS;
651 	next_tree = find_next_bit(data->pending_trees, CONNCOUNT_SLOTS, next_tree);
652 
653 	if (next_tree < CONNCOUNT_SLOTS) {
654 		data->gc_tree = next_tree;
655 		schedule_work(work);
656 	} else {
657 		data->gc_tree = 0;
658 	}
659 
660 	spin_unlock_bh(&root->lock);
661 }
662 
663 /* Count and return number of conntrack entries in 'net' with particular 'key'.
664  * If 'skb' is not null, insert the corresponding tuple into the accounting
665  * data structure. Call with RCU read lock.
666  */
nf_conncount_count_skb(struct net * net,const struct sk_buff * skb,u16 l3num,struct nf_conncount_data * data,const u32 * key)667 unsigned int nf_conncount_count_skb(struct net *net,
668 				    const struct sk_buff *skb,
669 				    u16 l3num,
670 				    struct nf_conncount_data *data,
671 				    const u32 *key)
672 {
673 	return count_tree(net, skb, l3num, data, key);
674 
675 }
676 EXPORT_SYMBOL_GPL(nf_conncount_count_skb);
677 
nf_conncount_root_init(struct nf_conncount_root * r)678 static void nf_conncount_root_init(struct nf_conncount_root *r)
679 {
680 	r->root = RB_ROOT;
681 	spin_lock_init(&r->lock);
682 	seqcount_spinlock_init(&r->count, &r->lock);
683 }
684 
nf_conncount_init(struct net * net,unsigned int keylen)685 struct nf_conncount_data *nf_conncount_init(struct net *net, unsigned int keylen)
686 {
687 	struct nf_conncount_data *data;
688 	int i;
689 
690 	if (keylen % sizeof(u32) ||
691 	    keylen / sizeof(u32) > MAX_KEYLEN ||
692 	    keylen == 0)
693 		return ERR_PTR(-EINVAL);
694 
695 	data = kvzalloc_obj(*data);
696 	if (!data)
697 		return ERR_PTR(-ENOMEM);
698 
699 	for (i = 0; i < ARRAY_SIZE(data->root); ++i)
700 		nf_conncount_root_init(&data->root[i]);
701 
702 	data->keylen = keylen / sizeof(u32);
703 	data->net = net;
704 	data->initval = get_random_u32();
705 	INIT_WORK(&data->gc_work, tree_gc_worker);
706 
707 	return data;
708 }
709 EXPORT_SYMBOL_GPL(nf_conncount_init);
710 
nf_conncount_cache_free(struct nf_conncount_list * list)711 void nf_conncount_cache_free(struct nf_conncount_list *list)
712 {
713 	struct nf_conncount_tuple *conn, *conn_n;
714 
715 	list_for_each_entry_safe(conn, conn_n, &list->head, node)
716 		kmem_cache_free(conncount_conn_cachep, conn);
717 }
718 EXPORT_SYMBOL_GPL(nf_conncount_cache_free);
719 
destroy_tree(struct nf_conncount_root * r)720 static void destroy_tree(struct nf_conncount_root *r)
721 {
722 	struct nf_conncount_rb *rbconn;
723 	struct rb_node *node;
724 
725 	while ((node = rb_first(&r->root)) != NULL) {
726 		rbconn = rb_entry(node, struct nf_conncount_rb, node);
727 
728 		rb_erase(node, &r->root);
729 
730 		nf_conncount_cache_free(&rbconn->list);
731 
732 		kmem_cache_free(conncount_rb_cachep, rbconn);
733 	}
734 }
735 
nf_conncount_destroy(struct net * net,struct nf_conncount_data * data)736 void nf_conncount_destroy(struct net *net, struct nf_conncount_data *data)
737 {
738 	unsigned int i;
739 
740 	disable_work_sync(&data->gc_work);
741 
742 	for (i = 0; i < ARRAY_SIZE(data->root); ++i)
743 		destroy_tree(&data->root[i]);
744 
745 	kvfree(data);
746 }
747 EXPORT_SYMBOL_GPL(nf_conncount_destroy);
748 
nf_conncount_modinit(void)749 static int __init nf_conncount_modinit(void)
750 {
751 	conncount_conn_cachep = KMEM_CACHE(nf_conncount_tuple, 0);
752 	if (!conncount_conn_cachep)
753 		return -ENOMEM;
754 
755 	conncount_rb_cachep = KMEM_CACHE(nf_conncount_rb, 0);
756 	if (!conncount_rb_cachep) {
757 		kmem_cache_destroy(conncount_conn_cachep);
758 		return -ENOMEM;
759 	}
760 
761 	return 0;
762 }
763 
nf_conncount_modexit(void)764 static void __exit nf_conncount_modexit(void)
765 {
766 	rcu_barrier();
767 	kmem_cache_destroy(conncount_conn_cachep);
768 	kmem_cache_destroy(conncount_rb_cachep);
769 }
770 
771 module_init(nf_conncount_modinit);
772 module_exit(nf_conncount_modexit);
773 MODULE_AUTHOR("Jan Engelhardt <jengelh@medozas.de>");
774 MODULE_AUTHOR("Florian Westphal <fw@strlen.de>");
775 MODULE_DESCRIPTION("netfilter: count number of connections matching a key");
776 MODULE_LICENSE("GPL");
777