net/netfilter/nf_conncount.c

   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 <[email protected]>
   7  *
   8  * split from xt_connlimit.c:
   9  *   (c) 2000 Gerd Knorr <[email protected]>
  10  *   Nov 2002: Martin Bene <[email protected]>:
  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 MAX_KEYLEN              5
  39
  40 /* we will save the tuples of all connections we care about */
  41 struct nf_conncount_tuple {
  42         struct list_head                node;
  43         struct nf_conntrack_tuple       tuple;
  44         struct nf_conntrack_zone        zone;
  45         int                             cpu;
  46         u32                             jiffies32;
  47 };
  48
  49 struct nf_conncount_rb {
  50         struct rb_node node;
  51         struct nf_conncount_list list;
  52         u32 key[MAX_KEYLEN];
  53         struct rcu_head rcu_head;
  54 };
  55
  56 static spinlock_t nf_conncount_locks[CONNCOUNT_SLOTS] __cacheline_aligned_in_smp;
  57
  58 struct nf_conncount_data {
  59         unsigned int keylen;
  60         struct rb_root root[CONNCOUNT_SLOTS];
  61         struct net *net;
  62         struct work_struct gc_work;
  63         unsigned long pending_trees[BITS_TO_LONGS(CONNCOUNT_SLOTS)];
  64         unsigned int gc_tree;
  65 };
  66
  67 static u_int32_t conncount_rnd __read_mostly;
  68 static struct kmem_cache *conncount_rb_cachep __read_mostly;
  69 static struct kmem_cache *conncount_conn_cachep __read_mostly;
  70
  71 static inline bool already_closed(const struct nf_conn *conn)
  72 {
  73         if (nf_ct_protonum(conn) == IPPROTO_TCP)
  74                 return conn->proto.tcp.state == TCP_CONNTRACK_TIME_WAIT ||
  75                        conn->proto.tcp.state == TCP_CONNTRACK_CLOSE;
  76         else
  77                 return false;
  78 }
  79
  80 static int key_diff(const u32 *a, const u32 *b, unsigned int klen)
  81 {
  82         return memcmp(a, b, klen * sizeof(u32));
  83 }
  84
  85 static void conn_free(struct nf_conncount_list *list,
  86                       struct nf_conncount_tuple *conn)
  87 {
  88         lockdep_assert_held(&list->list_lock);
  89
  90         list->count--;
  91         list_del(&conn->node);
  92
  93         kmem_cache_free(conncount_conn_cachep, conn);
  94 }
  95
  96 static const struct nf_conntrack_tuple_hash *
  97 find_or_evict(struct net *net, struct nf_conncount_list *list,
  98               struct nf_conncount_tuple *conn)
  99 {
 100         const struct nf_conntrack_tuple_hash *found;
 101         unsigned long a, b;
 102         int cpu = raw_smp_processor_id();
 103         u32 age;
 104
 105         found = nf_conntrack_find_get(net, &conn->zone, &conn->tuple);
 106         if (found)
 107                 return found;
 108         b = conn->jiffies32;
 109         a = (u32)jiffies;
 110
 111         /* conn might have been added just before by another cpu and
 112          * might still be unconfirmed.  In this case, nf_conntrack_find()
 113          * returns no result.  Thus only evict if this cpu added the
 114          * stale entry or if the entry is older than two jiffies.
 115          */
 116         age = a - b;
 117         if (conn->cpu == cpu || age >= 2) {
 118                 conn_free(list, conn);
 119                 return ERR_PTR(-ENOENT);
 120         }
 121
 122         return ERR_PTR(-EAGAIN);
 123 }
 124
 125 static int __nf_conncount_add(struct net *net,
 126                               struct nf_conncount_list *list,
 127                               const struct nf_conntrack_tuple *tuple,
 128                               const struct nf_conntrack_zone *zone)
 129 {
 130         const struct nf_conntrack_tuple_hash *found;
 131         struct nf_conncount_tuple *conn, *conn_n;
 132         struct nf_conn *found_ct;
 133         unsigned int collect = 0;
 134
 135         if (time_is_after_eq_jiffies((unsigned long)list->last_gc))
 136                 goto add_new_node;
 137
 138         /* check the saved connections */
 139         list_for_each_entry_safe(conn, conn_n, &list->head, node) {
 140                 if (collect > CONNCOUNT_GC_MAX_NODES)
 141                         break;
 142
 143                 found = find_or_evict(net, list, conn);
 144                 if (IS_ERR(found)) {
 145                         /* Not found, but might be about to be confirmed */
 146                         if (PTR_ERR(found) == -EAGAIN) {
 147                                 if (nf_ct_tuple_equal(&conn->tuple, tuple) &&
 148                                     nf_ct_zone_id(&conn->zone, conn->zone.dir) ==
 149                                     nf_ct_zone_id(zone, zone->dir))
 150                                         return 0; /* already exists */
 151                         } else {
 152                                 collect++;
 153                         }
 154                         continue;
 155                 }
 156
 157                 found_ct = nf_ct_tuplehash_to_ctrack(found);
 158
 159                 if (nf_ct_tuple_equal(&conn->tuple, tuple) &&
 160                     nf_ct_zone_equal(found_ct, zone, zone->dir)) {
 161                         /*
 162                          * We should not see tuples twice unless someone hooks
 163                          * this into a table without "-p tcp --syn".
 164                          *
 165                          * Attempt to avoid a re-add in this case.
 166                          */
 167                         nf_ct_put(found_ct);
 168                         return 0;
 169                 } else if (already_closed(found_ct)) {
 170                         /*
 171                          * we do not care about connections which are
 172                          * closed already -> ditch it
 173                          */
 174                         nf_ct_put(found_ct);
 175                         conn_free(list, conn);
 176                         collect++;
 177                         continue;
 178                 }
 179
 180                 nf_ct_put(found_ct);
 181         }
 182
 183 add_new_node:
 184         if (WARN_ON_ONCE(list->count > INT_MAX))
 185                 return -EOVERFLOW;
 186
 187         conn = kmem_cache_alloc(conncount_conn_cachep, GFP_ATOMIC);
 188         if (conn == NULL)
 189                 return -ENOMEM;
 190
 191         conn->tuple = *tuple;
 192         conn->zone = *zone;
 193         conn->cpu = raw_smp_processor_id();
 194         conn->jiffies32 = (u32)jiffies;
 195         list_add_tail(&conn->node, &list->head);
 196         list->count++;
 197         list->last_gc = (u32)jiffies;
 198         return 0;
 199 }
 200
 201 int nf_conncount_add(struct net *net,
 202                      struct nf_conncount_list *list,
 203                      const struct nf_conntrack_tuple *tuple,
 204                      const struct nf_conntrack_zone *zone)
 205 {
 206         int ret;
 207
 208         /* check the saved connections */
 209         spin_lock_bh(&list->list_lock);
 210         ret = __nf_conncount_add(net, list, tuple, zone);
 211         spin_unlock_bh(&list->list_lock);
 212
 213         return ret;
 214 }
 215 EXPORT_SYMBOL_GPL(nf_conncount_add);
 216
 217 void nf_conncount_list_init(struct nf_conncount_list *list)
 218 {
 219         spin_lock_init(&list->list_lock);
 220         INIT_LIST_HEAD(&list->head);
 221         list->count = 0;
 222         list->last_gc = (u32)jiffies;
 223 }
 224 EXPORT_SYMBOL_GPL(nf_conncount_list_init);
 225
 226 /* Return true if the list is empty. Must be called with BH disabled. */
 227 bool nf_conncount_gc_list(struct net *net,
 228                           struct nf_conncount_list *list)
 229 {
 230         const struct nf_conntrack_tuple_hash *found;
 231         struct nf_conncount_tuple *conn, *conn_n;
 232         struct nf_conn *found_ct;
 233         unsigned int collected = 0;
 234         bool ret = false;
 235
 236         /* don't bother if we just did GC */
 237         if (time_is_after_eq_jiffies((unsigned long)READ_ONCE(list->last_gc)))
 238                 return false;
 239
 240         /* don't bother if other cpu is already doing GC */
 241         if (!spin_trylock(&list->list_lock))
 242                 return false;
 243
 244         list_for_each_entry_safe(conn, conn_n, &list->head, node) {
 245                 found = find_or_evict(net, list, conn);
 246                 if (IS_ERR(found)) {
 247                         if (PTR_ERR(found) == -ENOENT)
 248                                 collected++;
 249                         continue;
 250                 }
 251
 252                 found_ct = nf_ct_tuplehash_to_ctrack(found);
 253                 if (already_closed(found_ct)) {
 254                         /*
 255                          * we do not care about connections which are
 256                          * closed already -> ditch it
 257                          */
 258                         nf_ct_put(found_ct);
 259                         conn_free(list, conn);
 260                         collected++;
 261                         continue;
 262                 }
 263
 264                 nf_ct_put(found_ct);
 265                 if (collected > CONNCOUNT_GC_MAX_NODES)
 266                         break;
 267         }
 268
 269         if (!list->count)
 270                 ret = true;
 271         list->last_gc = (u32)jiffies;
 272         spin_unlock(&list->list_lock);
 273
 274         return ret;
 275 }
 276 EXPORT_SYMBOL_GPL(nf_conncount_gc_list);
 277
 278 static void __tree_nodes_free(struct rcu_head *h)
 279 {
 280         struct nf_conncount_rb *rbconn;
 281
 282         rbconn = container_of(h, struct nf_conncount_rb, rcu_head);
 283         kmem_cache_free(conncount_rb_cachep, rbconn);
 284 }
 285
 286 /* caller must hold tree nf_conncount_locks[] lock */
 287 static void tree_nodes_free(struct rb_root *root,
 288                             struct nf_conncount_rb *gc_nodes[],
 289                             unsigned int gc_count)
 290 {
 291         struct nf_conncount_rb *rbconn;
 292
 293         while (gc_count) {
 294                 rbconn = gc_nodes[--gc_count];
 295                 spin_lock(&rbconn->list.list_lock);
 296                 if (!rbconn->list.count) {
 297                         rb_erase(&rbconn->node, root);
 298                         call_rcu(&rbconn->rcu_head, __tree_nodes_free);
 299                 }
 300                 spin_unlock(&rbconn->list.list_lock);
 301         }
 302 }
 303
 304 static void schedule_gc_worker(struct nf_conncount_data *data, int tree)
 305 {
 306         set_bit(tree, data->pending_trees);
 307         schedule_work(&data->gc_work);
 308 }
 309
 310 static unsigned int
 311 insert_tree(struct net *net,
 312             struct nf_conncount_data *data,
 313             struct rb_root *root,
 314             unsigned int hash,
 315             const u32 *key,
 316             const struct nf_conntrack_tuple *tuple,
 317             const struct nf_conntrack_zone *zone)
 318 {
 319         struct nf_conncount_rb *gc_nodes[CONNCOUNT_GC_MAX_NODES];
 320         struct rb_node **rbnode, *parent;
 321         struct nf_conncount_rb *rbconn;
 322         struct nf_conncount_tuple *conn;
 323         unsigned int count = 0, gc_count = 0;
 324         u8 keylen = data->keylen;
 325         bool do_gc = true;
 326
 327         spin_lock_bh(&nf_conncount_locks[hash]);
 328 restart:
 329         parent = NULL;
 330         rbnode = &(root->rb_node);
 331         while (*rbnode) {
 332                 int diff;
 333                 rbconn = rb_entry(*rbnode, struct nf_conncount_rb, node);
 334
 335                 parent = *rbnode;
 336                 diff = key_diff(key, rbconn->key, keylen);
 337                 if (diff < 0) {
 338                         rbnode = &((*rbnode)->rb_left);
 339                 } else if (diff > 0) {
 340                         rbnode = &((*rbnode)->rb_right);
 341                 } else {
 342                         int ret;
 343
 344                         ret = nf_conncount_add(net, &rbconn->list, tuple, zone);
 345                         if (ret)
 346                                 count = 0; /* hotdrop */
 347                         else
 348                                 count = rbconn->list.count;
 349                         tree_nodes_free(root, gc_nodes, gc_count);
 350                         goto out_unlock;
 351                 }
 352
 353                 if (gc_count >= ARRAY_SIZE(gc_nodes))
 354                         continue;
 355
 356                 if (do_gc && nf_conncount_gc_list(net, &rbconn->list))
 357                         gc_nodes[gc_count++] = rbconn;
 358         }
 359
 360         if (gc_count) {
 361                 tree_nodes_free(root, gc_nodes, gc_count);
 362                 schedule_gc_worker(data, hash);
 363                 gc_count = 0;
 364                 do_gc = false;
 365                 goto restart;
 366         }
 367
 368         /* expected case: match, insert new node */
 369         rbconn = kmem_cache_alloc(conncount_rb_cachep, GFP_ATOMIC);
 370         if (rbconn == NULL)
 371                 goto out_unlock;
 372
 373         conn = kmem_cache_alloc(conncount_conn_cachep, GFP_ATOMIC);
 374         if (conn == NULL) {
 375                 kmem_cache_free(conncount_rb_cachep, rbconn);
 376                 goto out_unlock;
 377         }
 378
 379         conn->tuple = *tuple;
 380         conn->zone = *zone;
 381         memcpy(rbconn->key, key, sizeof(u32) * keylen);
 382
 383         nf_conncount_list_init(&rbconn->list);
 384         list_add(&conn->node, &rbconn->list.head);
 385         count = 1;
 386         rbconn->list.count = count;
 387
 388         rb_link_node_rcu(&rbconn->node, parent, rbnode);
 389         rb_insert_color(&rbconn->node, root);
 390 out_unlock:
 391         spin_unlock_bh(&nf_conncount_locks[hash]);
 392         return count;
 393 }
 394
 395 static unsigned int
 396 count_tree(struct net *net,
 397            struct nf_conncount_data *data,
 398            const u32 *key,
 399            const struct nf_conntrack_tuple *tuple,
 400            const struct nf_conntrack_zone *zone)
 401 {
 402         struct rb_root *root;
 403         struct rb_node *parent;
 404         struct nf_conncount_rb *rbconn;
 405         unsigned int hash;
 406         u8 keylen = data->keylen;
 407
 408         hash = jhash2(key, data->keylen, conncount_rnd) % CONNCOUNT_SLOTS;
 409         root = &data->root[hash];
 410
 411         parent = rcu_dereference_raw(root->rb_node);
 412         while (parent) {
 413                 int diff;
 414
 415                 rbconn = rb_entry(parent, struct nf_conncount_rb, node);
 416
 417                 diff = key_diff(key, rbconn->key, keylen);
 418                 if (diff < 0) {
 419                         parent = rcu_dereference_raw(parent->rb_left);
 420                 } else if (diff > 0) {
 421                         parent = rcu_dereference_raw(parent->rb_right);
 422                 } else {
 423                         int ret;
 424
 425                         if (!tuple) {
 426                                 nf_conncount_gc_list(net, &rbconn->list);
 427                                 return rbconn->list.count;
 428                         }
 429
 430                         spin_lock_bh(&rbconn->list.list_lock);
 431                         /* Node might be about to be free'd.
 432                          * We need to defer to insert_tree() in this case.
 433                          */
 434                         if (rbconn->list.count == 0) {
 435                                 spin_unlock_bh(&rbconn->list.list_lock);
 436                                 break;
 437                         }
 438
 439                         /* same source network -> be counted! */
 440                         ret = __nf_conncount_add(net, &rbconn->list, tuple, zone);
 441                         spin_unlock_bh(&rbconn->list.list_lock);
 442                         if (ret)
 443                                 return 0; /* hotdrop */
 444                         else
 445                                 return rbconn->list.count;
 446                 }
 447         }
 448
 449         if (!tuple)
 450                 return 0;
 451
 452         return insert_tree(net, data, root, hash, key, tuple, zone);
 453 }
 454
 455 static void tree_gc_worker(struct work_struct *work)
 456 {
 457         struct nf_conncount_data *data = container_of(work, struct nf_conncount_data, gc_work);
 458         struct nf_conncount_rb *gc_nodes[CONNCOUNT_GC_MAX_NODES], *rbconn;
 459         struct rb_root *root;
 460         struct rb_node *node;
 461         unsigned int tree, next_tree, gc_count = 0;
 462
 463         tree = data->gc_tree % CONNCOUNT_SLOTS;
 464         root = &data->root[tree];
 465
 466         local_bh_disable();
 467         rcu_read_lock();
 468         for (node = rb_first(root); node != NULL; node = rb_next(node)) {
 469                 rbconn = rb_entry(node, struct nf_conncount_rb, node);
 470                 if (nf_conncount_gc_list(data->net, &rbconn->list))
 471                         gc_count++;
 472         }
 473         rcu_read_unlock();
 474         local_bh_enable();
 475
 476         cond_resched();
 477
 478         spin_lock_bh(&nf_conncount_locks[tree]);
 479         if (gc_count < ARRAY_SIZE(gc_nodes))
 480                 goto next; /* do not bother */
 481
 482         gc_count = 0;
 483         node = rb_first(root);
 484         while (node != NULL) {
 485                 rbconn = rb_entry(node, struct nf_conncount_rb, node);
 486                 node = rb_next(node);
 487
 488                 if (rbconn->list.count > 0)
 489                         continue;
 490
 491                 gc_nodes[gc_count++] = rbconn;
 492                 if (gc_count >= ARRAY_SIZE(gc_nodes)) {
 493                         tree_nodes_free(root, gc_nodes, gc_count);
 494                         gc_count = 0;
 495                 }
 496         }
 497
 498         tree_nodes_free(root, gc_nodes, gc_count);
 499 next:
 500         clear_bit(tree, data->pending_trees);
 501
 502         next_tree = (tree + 1) % CONNCOUNT_SLOTS;
 503         next_tree = find_next_bit(data->pending_trees, CONNCOUNT_SLOTS, next_tree);
 504
 505         if (next_tree < CONNCOUNT_SLOTS) {
 506                 data->gc_tree = next_tree;
 507                 schedule_work(work);
 508         }
 509
 510         spin_unlock_bh(&nf_conncount_locks[tree]);
 511 }
 512
 513 /* Count and return number of conntrack entries in 'net' with particular 'key'.
 514  * If 'tuple' is not null, insert it into the accounting data structure.
 515  * Call with RCU read lock.
 516  */
 517 unsigned int nf_conncount_count(struct net *net,
 518                                 struct nf_conncount_data *data,
 519                                 const u32 *key,
 520                                 const struct nf_conntrack_tuple *tuple,
 521                                 const struct nf_conntrack_zone *zone)
 522 {
 523         return count_tree(net, data, key, tuple, zone);
 524 }
 525 EXPORT_SYMBOL_GPL(nf_conncount_count);
 526
 527 struct nf_conncount_data *nf_conncount_init(struct net *net, unsigned int family,
 528                                             unsigned int keylen)
 529 {
 530         struct nf_conncount_data *data;
 531         int ret, i;
 532
 533         if (keylen % sizeof(u32) ||
 534             keylen / sizeof(u32) > MAX_KEYLEN ||
 535             keylen == 0)
 536                 return ERR_PTR(-EINVAL);
 537
 538         net_get_random_once(&conncount_rnd, sizeof(conncount_rnd));
 539
 540         data = kmalloc(sizeof(*data), GFP_KERNEL);
 541         if (!data)
 542                 return ERR_PTR(-ENOMEM);
 543
 544         ret = nf_ct_netns_get(net, family);
 545         if (ret < 0) {
 546                 kfree(data);
 547                 return ERR_PTR(ret);
 548         }
 549
 550         for (i = 0; i < ARRAY_SIZE(data->root); ++i)
 551                 data->root[i] = RB_ROOT;
 552
 553         data->keylen = keylen / sizeof(u32);
 554         data->net = net;
 555         INIT_WORK(&data->gc_work, tree_gc_worker);
 556
 557         return data;
 558 }
 559 EXPORT_SYMBOL_GPL(nf_conncount_init);
 560
 561 void nf_conncount_cache_free(struct nf_conncount_list *list)
 562 {
 563         struct nf_conncount_tuple *conn, *conn_n;
 564
 565         list_for_each_entry_safe(conn, conn_n, &list->head, node)
 566                 kmem_cache_free(conncount_conn_cachep, conn);
 567 }
 568 EXPORT_SYMBOL_GPL(nf_conncount_cache_free);
 569
 570 static void destroy_tree(struct rb_root *r)
 571 {
 572         struct nf_conncount_rb *rbconn;
 573         struct rb_node *node;
 574
 575         while ((node = rb_first(r)) != NULL) {
 576                 rbconn = rb_entry(node, struct nf_conncount_rb, node);
 577
 578                 rb_erase(node, r);
 579
 580                 nf_conncount_cache_free(&rbconn->list);
 581
 582                 kmem_cache_free(conncount_rb_cachep, rbconn);
 583         }
 584 }
 585
 586 void nf_conncount_destroy(struct net *net, unsigned int family,
 587                           struct nf_conncount_data *data)
 588 {
 589         unsigned int i;
 590
 591         cancel_work_sync(&data->gc_work);
 592         nf_ct_netns_put(net, family);
 593
 594         for (i = 0; i < ARRAY_SIZE(data->root); ++i)
 595                 destroy_tree(&data->root[i]);
 596
 597         kfree(data);
 598 }
 599 EXPORT_SYMBOL_GPL(nf_conncount_destroy);
 600
 601 static int __init nf_conncount_modinit(void)
 602 {
 603         int i;
 604
 605         for (i = 0; i < CONNCOUNT_SLOTS; ++i)
 606                 spin_lock_init(&nf_conncount_locks[i]);
 607
 608         conncount_conn_cachep = kmem_cache_create("nf_conncount_tuple",
 609                                            sizeof(struct nf_conncount_tuple),
 610                                            0, 0, NULL);
 611         if (!conncount_conn_cachep)
 612                 return -ENOMEM;
 613
 614         conncount_rb_cachep = kmem_cache_create("nf_conncount_rb",
 615                                            sizeof(struct nf_conncount_rb),
 616                                            0, 0, NULL);
 617         if (!conncount_rb_cachep) {
 618                 kmem_cache_destroy(conncount_conn_cachep);
 619                 return -ENOMEM;
 620         }
 621
 622         return 0;
 623 }
 624
 625 static void __exit nf_conncount_modexit(void)
 626 {
 627         kmem_cache_destroy(conncount_conn_cachep);
 628         kmem_cache_destroy(conncount_rb_cachep);
 629 }
 630
 631 module_init(nf_conncount_modinit);
 632 module_exit(nf_conncount_modexit);
 633 MODULE_AUTHOR("Jan Engelhardt <[email protected]>");
 634 MODULE_AUTHOR("Florian Westphal <[email protected]>");
 635 MODULE_DESCRIPTION("netfilter: count number of connections matching a key");
 636 MODULE_LICENSE("GPL");