use BPF_MAP_TYPE_LPM_TRIE map in XDP program to block a range

Signed-off-by: Y7n05h <Y7n05h@protonmail.com>

diff --git a/contrib/xdp-filter.ebpf.src b/contrib/xdp-filter.ebpf.src
index 5ae8a84020d4911383e80ae2b9a97823fdbc9f22..9b0867eaeadda9264e7dcb2c48b6c0cfcffc6d21 100644 (file)
--- a/contrib/xdp-filter.ebpf.src
+++ b/contrib/xdp-filter.ebpf.src
@@ -94,6 +94,17 @@ enum dns_action : uint8_t {
   TC = 2
 };
 
+struct CIDR4
+{
+  uint32_t cidr;
+  uint32_t addr;
+};
+struct CIDR6
+{
+  uint32_t cidr;
+  struct in6_addr addr;
+};
+
 /*
  * Store the matching counter and the associated action for a blocked element
  */
@@ -107,6 +118,18 @@ BPF_TABLE_PINNED("hash", uint32_t, struct map_value, v4filter, 1024, "/sys/fs/bp
 BPF_TABLE_PINNED("hash", struct in6_addr, struct map_value, v6filter, 1024, "/sys/fs/bpf/dnsdist/addr-v6");
 BPF_TABLE_PINNED("hash", struct dns_qname, struct map_value, qnamefilter, 1024, "/sys/fs/bpf/dnsdist/qnames");
 
+/*
+ * bcc has added BPF_TABLE_PINNED7 to the latest commit of the master branch, but it has not yet been released.
+ * https://github.com/iovisor/bcc/commit/fff25a8d4d445c6156b65aa8a4016ce0d78ab7fb
+ */
+#ifndef BPF_TABLE_PINNED7
+#define BPF_TABLE_PINNED7(_table_type, _key_type, _leaf_type, _name, _max_entries, _pinned, _flags) \
+  BPF_F_TABLE(_table_type ":" _pinned, _key_type, _leaf_type, _name, _max_entries, _flags)
+#endif
+
+BPF_TABLE_PINNED7("lpm_trie", struct CIDR4, struct map_value, cidr4filter, 1024, "/sys/fs/bpf/dnsdist/cidr4", BPF_F_NO_PREALLOC);
+BPF_TABLE_PINNED7("lpm_trie", struct CIDR6, struct map_value, cidr6filter, 1024, "/sys/fs/bpf/dnsdist/cidr6", BPF_F_NO_PREALLOC);
+
 /*
  * Initializer of a cursor pointer
  *  Copyright 2020, NLnet Labs, All rights reserved.
@@ -282,7 +305,7 @@ static inline enum dns_action check_qname(struct cursor *c)
  *         TC if (modified) message needs to be replied
  *         DROP if message needs to be blocked
  */
-static inline enum dns_action udp_dns_reply_v4(struct cursor *c, uint32_t key)
+static inline enum dns_action udp_dns_reply_v4(struct cursor *c, struct CIDR4 *key)
 {
   struct udphdr  *udp;
   struct dnshdr  *dns;
@@ -297,7 +320,7 @@ static inline enum dns_action udp_dns_reply_v4(struct cursor *c, uint32_t key)
   }    
 
   // if the address is blocked, perform the corresponding action
-  struct map_value* value = v4filter.lookup(&key);
+  struct map_value* value = v4filter.lookup(&key->addr);
 
   if (value) {
     __sync_fetch_and_add(&value->counter, 1);
@@ -306,16 +329,26 @@ static inline enum dns_action udp_dns_reply_v4(struct cursor *c, uint32_t key)
     } else {
       return value->action;
     }
-  } else {
-    enum dns_action action = check_qname(c);
-       if (action == TC) {
-               return set_tc_bit(udp, dns);
-       } else {
-      return action;
+  }
+
+  key->cidr = 32;
+  key->addr=bpf_htonl(key->addr);
+  value = cidr4filter.lookup(key);
+  if (value) {
+    __sync_fetch_and_add(&value->counter, 1);
+    if (value->action == TC) {
+      return set_tc_bit(udp, dns);
+    }
+    else {
+      return value->action;
     }
   }
 
-  return PASS;
+  enum dns_action action = check_qname(c);
+  if (action == TC) {
+    return set_tc_bit(udp, dns);
+  }
+  return action;
 }
 
 /*
@@ -324,7 +357,7 @@ static inline enum dns_action udp_dns_reply_v4(struct cursor *c, uint32_t key)
  *         TC if (modified) message needs to be replied
  *         DROP if message needs to be blocked
  */
-static inline enum dns_action udp_dns_reply_v6(struct cursor *c, struct in6_addr key)
+static inline enum dns_action udp_dns_reply_v6(struct cursor *c, struct CIDR6* key)
 {
    struct udphdr  *udp;
    struct dnshdr  *dns;
@@ -341,7 +374,7 @@ static inline enum dns_action udp_dns_reply_v6(struct cursor *c, struct in6_addr
   }
 
   // if the address is blocked, perform the corresponding action
-  struct map_value* value = v6filter.lookup(&key);
+  struct map_value* value = v6filter.lookup(&key->addr);
 
   if (value) {
     __sync_fetch_and_add(&value->counter, 1);
@@ -350,16 +383,25 @@ static inline enum dns_action udp_dns_reply_v6(struct cursor *c, struct in6_addr
     } else {
       return value->action;
     }
-  } else {
-    enum dns_action action = check_qname(c);
-       if (action == TC) {
-               return set_tc_bit(udp, dns);
-       } else {
-      return action;
+  }
+
+  key->cidr = 128;
+  value = cidr6filter.lookup(key);
+  if (value) {
+    __sync_fetch_and_add(&value->counter, 1);
+    if (value->action == TC) {
+      return set_tc_bit(udp, dns);
+    }
+    else {
+      return value->action;
     }
   }
 
-  return PASS;
+  enum dns_action action = check_qname(c);
+  if (action == TC) {
+    return set_tc_bit(udp, dns);
+  }
+  return action;
 }
 
 int xdp_dns_filter(struct xdp_md *ctx)
@@ -383,8 +425,10 @@ int xdp_dns_filter(struct xdp_md *ctx)
                if (!(ipv4 = parse_iphdr(&c)) || bpf_htons(ipv4->protocol != IPPROTO_UDP)) {
                        return XDP_PASS;
                }
+      struct CIDR4 key;
+      key.addr = bpf_htonl(ipv4->saddr);
       // if TC bit must not be set, apply the action
-               if ((r = udp_dns_reply_v4(&c, bpf_htonl(ipv4->saddr))) != TC) {
+               if ((r = udp_dns_reply_v4(&c, &key)) != TC) {
            return r == DROP ? XDP_DROP : XDP_PASS;
                }
   
@@ -400,8 +444,10 @@ int xdp_dns_filter(struct xdp_md *ctx)
                if (!(ipv6 = parse_ipv6hdr(&c)) || bpf_htons(ipv6->nexthdr != IPPROTO_UDP)) {
                        return XDP_PASS;
                } 
+      struct CIDR6 key;
+      key.addr = ipv6->saddr;
       // if TC bit must not be set, apply the action
-               if ((r = udp_dns_reply_v6(&c, ipv6->saddr)) != TC) {
+               if ((r = udp_dns_reply_v6(&c, &key)) != TC) {
                        return r == DROP ? XDP_DROP : XDP_PASS;
                }
   

diff --git a/contrib/xdp.py b/contrib/xdp.py
index c84893ae8e116f3442f575b7a2cf72ac0f394b0d..7391c5b72187c7edfa2e5d6cad7b561eefd6380a 100644 (file)
--- a/contrib/xdp.py
+++ b/contrib/xdp.py
@@ -18,9 +18,12 @@ DEV = "eth0"
 
 # The list of blocked IPv4, IPv6 and QNames
 # IP format : (IPAddress, Action)
+# CIDR format : (IPAddress/cidr, Action)
 # QName format : (QName, QType, Action)
 blocked_ipv4 = [("192.0.2.1", TC_ACTION)]
 blocked_ipv6 = [("2001:db8::1", TC_ACTION)]
+blocked_cidr4 = [("192.0.1.1/24", TC_ACTION)]
+blocked_cidr6 = [("2001:db8::1/128", TC_ACTION)]
 blocked_qnames = [("localhost", "A", DROP_ACTION), ("test.com", "*", TC_ACTION)]
 
 # Main
@@ -31,6 +34,8 @@ xdp.attach_xdp(DEV, fn, 0)
 
 v4filter = xdp.get_table("v4filter")
 v6filter = xdp.get_table("v6filter")
+cidr4filter = xdp.get_table("cidr4filter")
+cidr6filter = xdp.get_table("cidr6filter")
 qnamefilter = xdp.get_table("qnamefilter")
 
 for ip in blocked_ipv4:
@@ -51,6 +56,42 @@ for ip in blocked_ipv6:
   leaf.action = ip[1]
   v6filter[key] = leaf
 
+for item in blocked_cidr4:
+  print(f"Blocking {item}")
+  key = cidr4filter.Key()
+  cidr=32
+  tmp=item[0].split('/')
+  if len(tmp)>1:
+    cidr=int(tmp[1])
+    if cidr>32 or cidr<=0:
+      raise RuntimeError("params invalid")
+  ip=tmp[0]
+  key.cidr=cidr
+  key.addr = int.from_bytes(bytes(netaddr.IPAddress(ip)), "little", signed=False)
+  leaf = cidr4filter.Leaf()
+  leaf.counter = 0
+  leaf.action = item[1] 
+  cidr4filter[key] = leaf
+
+for item in blocked_cidr6:
+  print(f"Blocking {item}")
+  key = cidr6filter.Key()
+  cidr=128
+  tmp=item[0].split('/')
+  if len(tmp)>1:
+    cidr=int(tmp[1])
+    if cidr>128 or cidr<=0:
+      raise RuntimeError("params invalid")
+  ip=tmp[0]
+  key.cidr=cidr
+  ipv6_int = int(netaddr.IPAddress(ip[0]).value)
+  ipv6_bytes = bytearray([(ipv6_int & (255 << 8*(15-i))) >> (8*(15-i)) for i in range(16)])
+  key = (ct.c_uint8 * 16).from_buffer(ipv6_bytes)
+  leaf = cidr6filter.Leaf()
+  leaf.counter = 0
+  leaf.action = item[1]
+  cidr6filter[key] = leaf
+
 for qname in blocked_qnames:
   print(f"Blocking {qname}")
   key = qnamefilter.Key()
@@ -77,6 +118,11 @@ for item in v4filter.items():
   print(f"{str(netaddr.IPAddress(item[0].value))} ({ACTIONS[item[1].action]}): {item[1].counter}")
 for item in v6filter.items():
   print(f"{str(socket.inet_ntop(socket.AF_INET6, item[0]))} ({ACTIONS[item[1].action]}): {item[1].counter}")
+for item in cidr4filter.items():
+  addr = int.from_bytes(netaddr.IPAddress(item[0].addr), "big", signed=False)
+  print(f"{str(addr)}/{str(item[0].cidr)} ({ACTIONS[item[1].action]}): {item[1].counter}")
+for item in cidr6filter.items():
+  print(f"{str(socket.inet_ntop(socket.AF_INET6, item[0].addr))}/{str(item[0].cidr)} ({ACTIONS[item[1].action]}): {item[1].counter}")
 for item in qnamefilter.items():
   print(f"{''.join(map(chr, item[0].qname)).strip()}/{INV_QTYPES[item[0].qtype]} ({ACTIONS[item[1].action]}): {item[1].counter}")