From: Matt Corallo <git@bluematt.me>
Date: Wed, 7 Apr 2021 19:36:40 +0000 (-0400)
Subject: Track both packet count and packet sizes in drop counts.
X-Git-Url: http://git.bitcoin.ninja/index.cgi?p=flowspec-xdp;a=commitdiff_plain;h=2400b2559ac2947300b90173a3fe25d632dba523

Track both packet count and packet sizes in drop counts.
---

diff --git a/dropcount.sh b/dropcount.sh
index 2620b5d..541dd30 100755
--- a/dropcount.sh
+++ b/dropcount.sh
@@ -1,17 +1,17 @@
 #!/bin/bash
 function PRINTCNT() {
-if [ "$KEY" != "" ]; then
-	if [ "$KEY" = "0" ]; then
-		echo -e "$CNT:\tInvalid packet length"
-	elif [ "$KEY" = "1" ]; then
-		echo -e "$CNT:\tInvalid VLAN tag"
-	elif [ "$KEY" = "2" ]; then
-		echo -e "$CNT:\tInvalid/rejected IHL IPv4 field"
-	elif [ "$KEY" = "3" ]; then
-		echo -e "$CNT:\tRejected IPv6 fragments"
+if [ "$1" != "" ]; then
+	if [ "$1" = "0" ]; then
+		echo -e "$2\t$3\tInvalid packet length"
+	elif [ "$1" = "1" ]; then
+		echo -e "$2\t$3\tInvalid VLAN tag"
+	elif [ "$1" = "2" ]; then
+		echo -e "$2\t$3\tInvalid/rejected IHL IPv4 field"
+	elif [ "$1" = "3" ]; then
+		echo -e "$2\t$3\tRejected IPv6 fragments"
 	else
-		echo -en "$CNT:\t"
-		cat "$(dirname ${BASH_SOURCE[0]})/installed-rules.txt" | head -n $(( $KEY - 3 )) | tail -n1
+		echo -en "$2\t$3\t"
+		cat "$(dirname ${BASH_SOURCE[0]})/installed-rules.txt" | head -n $(( $1 - 3 )) | tail -n1
 	fi
 fi
 CNT=0
@@ -20,26 +20,28 @@ MAP_CONTENTS="$(bpftool map show | grep drop_cnt_map | awk '{ print $1 }' | tr -
 	bpftool map dump id "$IF"
 done)"
 echo "$MAP_CONTENTS" | {
-	declare -a COUNTS
+	declare -a BYTES
+	declare -a PACKETS
 	KEY=""
 	while read LINE; do
 		case "$LINE" in
-			"key:") ;;
-			"value"*)
-				COUNTS["$KEY"]=$(( ${COUNTS["$KEY"]} + $(echo "$LINE" | awk '{ print "0x" $11 $10 $9 $8 $7 $6 $5 $4 }') ))
-				;;
-			"Found "*) ;;
-			*)
-				KEY=$((16#$(echo "$LINE" | awk '{ print $4 $3 $2 $1 }')))
-				if [ "$COUNTS["$KEY"]" = "" ]; then
-					COUNTS["$KEY"]=0
+			*"key"*)
+				KEY=$(echo "$LINE" | awk '{ print $2 }' | tr -d ',')
+				if [ "${BYTES["${KEY}"]}" = "" ]; then
+					BYTES["${KEY}"]=0
+					PACKETS["${KEY}"]=0
 				fi
 				;;
+			*"bytes"*)
+				BYTES["${KEY}"]=$(( ${BYTES["$KEY"]} + $(echo "$LINE" | awk '{ print $2 }' | tr -d ',') ))
+				;;
+			*"packets"*)
+				PACKETS["$KEY"]=$(( ${PACKETS["$KEY"]} + $(echo "$LINE" | awk '{ print $2 }' | tr -d ',') ))
+				;;
 		esac
 	done
-	for C in "${!COUNTS[@]}"; do
-		KEY=$C
-		CNT="${COUNTS["$KEY"]}"
-		PRINTCNT
+	echo -e "pkts\tKBytes\tRule"
+	for C in "${!BYTES[@]}"; do
+		PRINTCNT $C "${PACKETS["$C"]}" "$(( ${BYTES["$C"]} / 1000 ))"
 	done
 }
diff --git a/xdp.c b/xdp.c
index 8003b89..506d422 100644
--- a/xdp.c
+++ b/xdp.c
@@ -144,22 +144,29 @@ static const int XDP_PASS = 0;
 static const int XDP_DROP = 1;
 
 static long drop_cnt_map[RULECNT + STATIC_RULE_CNT];
-#define INCREMENT_MATCH(reason) drop_cnt_map[reason] += 1;
+#define INCREMENT_MATCH(reason) { drop_cnt_map[reason] += 1; drop_cnt_map[reason] += data_end - pktdata; }
 
 #else
 #include <linux/bpf.h>
 #include <bpf/bpf_helpers.h>
 
-struct bpf_map_def SEC("maps") drop_cnt_map = {
-	.type = BPF_MAP_TYPE_PERCPU_ARRAY,
-	.key_size = sizeof(uint32_t),
-	.value_size = sizeof(long),
-	.max_entries = RULECNT + STATIC_RULE_CNT,
+struct match_counter {
+	uint64_t bytes;
+	uint64_t packets;
 };
+struct {
+	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
+	__uint(max_entries, RULECNT + STATIC_RULE_CNT);
+	__u32 *key;
+	struct match_counter *value;
+} drop_cnt_map SEC(".maps");
+
 #define INCREMENT_MATCH(reason) { \
-	long *value = bpf_map_lookup_elem(&drop_cnt_map, &reason); \
-	if (value) \
-		*value += 1; \
+	struct match_counter *value = bpf_map_lookup_elem(&drop_cnt_map, &reason); \
+	if (value) { \
+		value->bytes += data_end - pktdata; \
+		value->packets += 1; \
+	} \
 }
 
 #ifdef RATE_CNT
@@ -186,30 +193,32 @@ int xdp_drop_prog(struct xdp_md *ctx)
 	unsigned short eth_proto;
 
 	{
+		// DO_RETURN in CHECK_LEN relies on pktdata being set to calculate packet length.
+		// That said, we don't want to overflow, so just set packet length to 0 here.
+		pktdata = data_end;
 		CHECK_LEN((size_t)ctx->data, ethhdr);
 		const struct ethhdr *const eth = (void*)(size_t)ctx->data;
+		pktdata = (const void *)(long)ctx->data + sizeof(struct ethhdr);
 
 #if PARSE_8021Q == PARSE
 		if (likely(eth->h_proto == BE16(ETH_P_8021Q))) {
 			CHECK_LEN((size_t)ctx->data, ethhdr_vlan);
 			const struct ethhdr_vlan *const eth_vlan = (void*)(size_t)ctx->data;
-
+			pktdata = (const void *)(long)ctx->data + sizeof(struct ethhdr_vlan);
 #ifdef REQ_8021Q
 			if (unlikely((eth_vlan->tci & BE16(0xfff)) != BE16(REQ_8021Q)))
 				DO_RETURN(VLAN_DROP, XDP_DROP);
 #endif
-
 			eth_proto = eth_vlan->h_proto;
-			pktdata = (const void *)(long)ctx->data + sizeof(struct ethhdr_vlan);
 #else
 		if (unlikely(eth->h_proto == BE16(ETH_P_8021Q))) {
+			pktdata = (const void *)(long)ctx->data + sizeof(struct ethhdr_vlan);
 			DO_RETURN(VLAN_DROP, PARSE_8021Q);
 #endif
 		} else {
 #ifdef REQ_8021Q
 			DO_RETURN(VLAN_DROP, XDP_DROP);
 #else
-			pktdata = (const void *)(long)ctx->data + sizeof(struct ethhdr);
 			eth_proto = eth->h_proto;
 #endif
 		}