net/wireguard/allowedips.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 /*
0003  * Copyright (C) 2015-2019 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
0004  */
0005
0006 #include "allowedips.h"
0007 #include "peer.h"
0008
0009 enum { MAX_ALLOWEDIPS_BITS = 128 };
0010
0011 static struct kmem_cache *node_cache;
0012
0013 static void swap_endian(u8 *dst, const u8 *src, u8 bits)
0014 {
0015     if (bits == 32) {
0016         *(u32 *)dst = be32_to_cpu(*(const __be32 *)src);
0017     } else if (bits == 128) {
0018         ((u64 *)dst)[0] = be64_to_cpu(((const __be64 *)src)[0]);
0019         ((u64 *)dst)[1] = be64_to_cpu(((const __be64 *)src)[1]);
0020     }
0021 }
0022
0023 static void copy_and_assign_cidr(struct allowedips_node *node, const u8 *src,
0024                  u8 cidr, u8 bits)
0025 {
0026     node->cidr = cidr;
0027     node->bit_at_a = cidr / 8U;
0028 #ifdef __LITTLE_ENDIAN
0029     node->bit_at_a ^= (bits / 8U - 1U) % 8U;
0030 #endif
0031     node->bit_at_b = 7U - (cidr % 8U);
0032     node->bitlen = bits;
0033     memcpy(node->bits, src, bits / 8U);
0034 }
0035
0036 static inline u8 choose(struct allowedips_node *node, const u8 *key)
0037 {
0038     return (key[node->bit_at_a] >> node->bit_at_b) & 1;
0039 }
0040
0041 static void push_rcu(struct allowedips_node **stack,
0042              struct allowedips_node __rcu *p, unsigned int *len)
0043 {
0044     if (rcu_access_pointer(p)) {
0045         if (WARN_ON(IS_ENABLED(DEBUG) && *len >= MAX_ALLOWEDIPS_BITS))
0046             return;
0047         stack[(*len)++] = rcu_dereference_raw(p);
0048     }
0049 }
0050
0051 static void node_free_rcu(struct rcu_head *rcu)
0052 {
0053     kmem_cache_free(node_cache, container_of(rcu, struct allowedips_node, rcu));
0054 }
0055
0056 static void root_free_rcu(struct rcu_head *rcu)
0057 {
0058     struct allowedips_node *node, *stack[MAX_ALLOWEDIPS_BITS] = {
0059         container_of(rcu, struct allowedips_node, rcu) };
0060     unsigned int len = 1;
0061
0062     while (len > 0 && (node = stack[--len])) {
0063         push_rcu(stack, node->bit[0], &len);
0064         push_rcu(stack, node->bit[1], &len);
0065         kmem_cache_free(node_cache, node);
0066     }
0067 }
0068
0069 static void root_remove_peer_lists(struct allowedips_node *root)
0070 {
0071     struct allowedips_node *node, *stack[MAX_ALLOWEDIPS_BITS] = { root };
0072     unsigned int len = 1;
0073
0074     while (len > 0 && (node = stack[--len])) {
0075         push_rcu(stack, node->bit[0], &len);
0076         push_rcu(stack, node->bit[1], &len);
0077         if (rcu_access_pointer(node->peer))
0078             list_del(&node->peer_list);
0079     }
0080 }
0081
0082 static unsigned int fls128(u64 a, u64 b)
0083 {
0084     return a ? fls64(a) + 64U : fls64(b);
0085 }
0086
0087 static u8 common_bits(const struct allowedips_node *node, const u8 *key,
0088               u8 bits)
0089 {
0090     if (bits == 32)
0091         return 32U - fls(*(const u32 *)node->bits ^ *(const u32 *)key);
0092     else if (bits == 128)
0093         return 128U - fls128(
0094             *(const u64 *)&node->bits[0] ^ *(const u64 *)&key[0],
0095             *(const u64 *)&node->bits[8] ^ *(const u64 *)&key[8]);
0096     return 0;
0097 }
0098
0099 static bool prefix_matches(const struct allowedips_node *node, const u8 *key,
0100                u8 bits)
0101 {
0102     /* This could be much faster if it actually just compared the common
0103      * bits properly, by precomputing a mask bswap(~0 << (32 - cidr)), and
0104      * the rest, but it turns out that common_bits is already super fast on
0105      * modern processors, even taking into account the unfortunate bswap.
0106      * So, we just inline it like this instead.
0107      */
0108     return common_bits(node, key, bits) >= node->cidr;
0109 }
0110
0111 static struct allowedips_node *find_node(struct allowedips_node *trie, u8 bits,
0112                      const u8 *key)
0113 {
0114     struct allowedips_node *node = trie, *found = NULL;
0115
0116     while (node && prefix_matches(node, key, bits)) {
0117         if (rcu_access_pointer(node->peer))
0118             found = node;
0119         if (node->cidr == bits)
0120             break;
0121         node = rcu_dereference_bh(node->bit[choose(node, key)]);
0122     }
0123     return found;
0124 }
0125
0126 /* Returns a strong reference to a peer */
0127 static struct wg_peer *lookup(struct allowedips_node __rcu *root, u8 bits,
0128                   const void *be_ip)
0129 {
0130     /* Aligned so it can be passed to fls/fls64 */
0131     u8 ip[16] __aligned(__alignof(u64));
0132     struct allowedips_node *node;
0133     struct wg_peer *peer = NULL;
0134
0135     swap_endian(ip, be_ip, bits);
0136
0137     rcu_read_lock_bh();
0138 retry:
0139     node = find_node(rcu_dereference_bh(root), bits, ip);
0140     if (node) {
0141         peer = wg_peer_get_maybe_zero(rcu_dereference_bh(node->peer));
0142         if (!peer)
0143             goto retry;
0144     }
0145     rcu_read_unlock_bh();
0146     return peer;
0147 }
0148
0149 static bool node_placement(struct allowedips_node __rcu *trie, const u8 *key,
0150                u8 cidr, u8 bits, struct allowedips_node **rnode,
0151                struct mutex *lock)
0152 {
0153     struct allowedips_node *node = rcu_dereference_protected(trie, lockdep_is_held(lock));
0154     struct allowedips_node *parent = NULL;
0155     bool exact = false;
0156
0157     while (node && node->cidr <= cidr && prefix_matches(node, key, bits)) {
0158         parent = node;
0159         if (parent->cidr == cidr) {
0160             exact = true;
0161             break;
0162         }
0163         node = rcu_dereference_protected(parent->bit[choose(parent, key)], lockdep_is_held(lock));
0164     }
0165     *rnode = parent;
0166     return exact;
0167 }
0168
0169 static inline void connect_node(struct allowedips_node __rcu **parent, u8 bit, struct allowedips_node *node)
0170 {
0171     node->parent_bit_packed = (unsigned long)parent | bit;
0172     rcu_assign_pointer(*parent, node);
0173 }
0174
0175 static inline void choose_and_connect_node(struct allowedips_node *parent, struct allowedips_node *node)
0176 {
0177     u8 bit = choose(parent, node->bits);
0178     connect_node(&parent->bit[bit], bit, node);
0179 }
0180
0181 static int add(struct allowedips_node __rcu **trie, u8 bits, const u8 *key,
0182            u8 cidr, struct wg_peer *peer, struct mutex *lock)
0183 {
0184     struct allowedips_node *node, *parent, *down, *newnode;
0185
0186     if (unlikely(cidr > bits || !peer))
0187         return -EINVAL;
0188
0189     if (!rcu_access_pointer(*trie)) {
0190         node = kmem_cache_zalloc(node_cache, GFP_KERNEL);
0191         if (unlikely(!node))
0192             return -ENOMEM;
0193         RCU_INIT_POINTER(node->peer, peer);
0194         list_add_tail(&node->peer_list, &peer->allowedips_list);
0195         copy_and_assign_cidr(node, key, cidr, bits);
0196         connect_node(trie, 2, node);
0197         return 0;
0198     }
0199     if (node_placement(*trie, key, cidr, bits, &node, lock)) {
0200         rcu_assign_pointer(node->peer, peer);
0201         list_move_tail(&node->peer_list, &peer->allowedips_list);
0202         return 0;
0203     }
0204
0205     newnode = kmem_cache_zalloc(node_cache, GFP_KERNEL);
0206     if (unlikely(!newnode))
0207         return -ENOMEM;
0208     RCU_INIT_POINTER(newnode->peer, peer);
0209     list_add_tail(&newnode->peer_list, &peer->allowedips_list);
0210     copy_and_assign_cidr(newnode, key, cidr, bits);
0211
0212     if (!node) {
0213         down = rcu_dereference_protected(*trie, lockdep_is_held(lock));
0214     } else {
0215         const u8 bit = choose(node, key);
0216         down = rcu_dereference_protected(node->bit[bit], lockdep_is_held(lock));
0217         if (!down) {
0218             connect_node(&node->bit[bit], bit, newnode);
0219             return 0;
0220         }
0221     }
0222     cidr = min(cidr, common_bits(down, key, bits));
0223     parent = node;
0224
0225     if (newnode->cidr == cidr) {
0226         choose_and_connect_node(newnode, down);
0227         if (!parent)
0228             connect_node(trie, 2, newnode);
0229         else
0230             choose_and_connect_node(parent, newnode);
0231         return 0;
0232     }
0233
0234     node = kmem_cache_zalloc(node_cache, GFP_KERNEL);
0235     if (unlikely(!node)) {
0236         list_del(&newnode->peer_list);
0237         kmem_cache_free(node_cache, newnode);
0238         return -ENOMEM;
0239     }
0240     INIT_LIST_HEAD(&node->peer_list);
0241     copy_and_assign_cidr(node, newnode->bits, cidr, bits);
0242
0243     choose_and_connect_node(node, down);
0244     choose_and_connect_node(node, newnode);
0245     if (!parent)
0246         connect_node(trie, 2, node);
0247     else
0248         choose_and_connect_node(parent, node);
0249     return 0;
0250 }
0251
0252 void wg_allowedips_init(struct allowedips *table)
0253 {
0254     table->root4 = table->root6 = NULL;
0255     table->seq = 1;
0256 }
0257
0258 void wg_allowedips_free(struct allowedips *table, struct mutex *lock)
0259 {
0260     struct allowedips_node __rcu *old4 = table->root4, *old6 = table->root6;
0261
0262     ++table->seq;
0263     RCU_INIT_POINTER(table->root4, NULL);
0264     RCU_INIT_POINTER(table->root6, NULL);
0265     if (rcu_access_pointer(old4)) {
0266         struct allowedips_node *node = rcu_dereference_protected(old4,
0267                             lockdep_is_held(lock));
0268
0269         root_remove_peer_lists(node);
0270         call_rcu(&node->rcu, root_free_rcu);
0271     }
0272     if (rcu_access_pointer(old6)) {
0273         struct allowedips_node *node = rcu_dereference_protected(old6,
0274                             lockdep_is_held(lock));
0275
0276         root_remove_peer_lists(node);
0277         call_rcu(&node->rcu, root_free_rcu);
0278     }
0279 }
0280
0281 int wg_allowedips_insert_v4(struct allowedips *table, const struct in_addr *ip,
0282                 u8 cidr, struct wg_peer *peer, struct mutex *lock)
0283 {
0284     /* Aligned so it can be passed to fls */
0285     u8 key[4] __aligned(__alignof(u32));
0286
0287     ++table->seq;
0288     swap_endian(key, (const u8 *)ip, 32);
0289     return add(&table->root4, 32, key, cidr, peer, lock);
0290 }
0291
0292 int wg_allowedips_insert_v6(struct allowedips *table, const struct in6_addr *ip,
0293                 u8 cidr, struct wg_peer *peer, struct mutex *lock)
0294 {
0295     /* Aligned so it can be passed to fls64 */
0296     u8 key[16] __aligned(__alignof(u64));
0297
0298     ++table->seq;
0299     swap_endian(key, (const u8 *)ip, 128);
0300     return add(&table->root6, 128, key, cidr, peer, lock);
0301 }
0302
0303 void wg_allowedips_remove_by_peer(struct allowedips *table,
0304                   struct wg_peer *peer, struct mutex *lock)
0305 {
0306     struct allowedips_node *node, *child, **parent_bit, *parent, *tmp;
0307     bool free_parent;
0308
0309     if (list_empty(&peer->allowedips_list))
0310         return;
0311     ++table->seq;
0312     list_for_each_entry_safe(node, tmp, &peer->allowedips_list, peer_list) {
0313         list_del_init(&node->peer_list);
0314         RCU_INIT_POINTER(node->peer, NULL);
0315         if (node->bit[0] && node->bit[1])
0316             continue;
0317         child = rcu_dereference_protected(node->bit[!rcu_access_pointer(node->bit[0])],
0318                           lockdep_is_held(lock));
0319         if (child)
0320             child->parent_bit_packed = node->parent_bit_packed;
0321         parent_bit = (struct allowedips_node **)(node->parent_bit_packed & ~3UL);
0322         *parent_bit = child;
0323         parent = (void *)parent_bit -
0324              offsetof(struct allowedips_node, bit[node->parent_bit_packed & 1]);
0325         free_parent = !rcu_access_pointer(node->bit[0]) &&
0326                   !rcu_access_pointer(node->bit[1]) &&
0327                   (node->parent_bit_packed & 3) <= 1 &&
0328                   !rcu_access_pointer(parent->peer);
0329         if (free_parent)
0330             child = rcu_dereference_protected(
0331                     parent->bit[!(node->parent_bit_packed & 1)],
0332                     lockdep_is_held(lock));
0333         call_rcu(&node->rcu, node_free_rcu);
0334         if (!free_parent)
0335             continue;
0336         if (child)
0337             child->parent_bit_packed = parent->parent_bit_packed;
0338         *(struct allowedips_node **)(parent->parent_bit_packed & ~3UL) = child;
0339         call_rcu(&parent->rcu, node_free_rcu);
0340     }
0341 }
0342
0343 int wg_allowedips_read_node(struct allowedips_node *node, u8 ip[16], u8 *cidr)
0344 {
0345     const unsigned int cidr_bytes = DIV_ROUND_UP(node->cidr, 8U);
0346     swap_endian(ip, node->bits, node->bitlen);
0347     memset(ip + cidr_bytes, 0, node->bitlen / 8U - cidr_bytes);
0348     if (node->cidr)
0349         ip[cidr_bytes - 1U] &= ~0U << (-node->cidr % 8U);
0350
0351     *cidr = node->cidr;
0352     return node->bitlen == 32 ? AF_INET : AF_INET6;
0353 }
0354
0355 /* Returns a strong reference to a peer */
0356 struct wg_peer *wg_allowedips_lookup_dst(struct allowedips *table,
0357                      struct sk_buff *skb)
0358 {
0359     if (skb->protocol == htons(ETH_P_IP))
0360         return lookup(table->root4, 32, &ip_hdr(skb)->daddr);
0361     else if (skb->protocol == htons(ETH_P_IPV6))
0362         return lookup(table->root6, 128, &ipv6_hdr(skb)->daddr);
0363     return NULL;
0364 }
0365
0366 /* Returns a strong reference to a peer */
0367 struct wg_peer *wg_allowedips_lookup_src(struct allowedips *table,
0368                      struct sk_buff *skb)
0369 {
0370     if (skb->protocol == htons(ETH_P_IP))
0371         return lookup(table->root4, 32, &ip_hdr(skb)->saddr);
0372     else if (skb->protocol == htons(ETH_P_IPV6))
0373         return lookup(table->root6, 128, &ipv6_hdr(skb)->saddr);
0374     return NULL;
0375 }
0376
0377 int __init wg_allowedips_slab_init(void)
0378 {
0379     node_cache = KMEM_CACHE(allowedips_node, 0);
0380     return node_cache ? 0 : -ENOMEM;
0381 }
0382
0383 void wg_allowedips_slab_uninit(void)
0384 {
0385     rcu_barrier();
0386     kmem_cache_destroy(node_cache);
0387 }
0388
0389 #include "selftest/allowedips.c"