From 77af5406a3a9bd1b84ebd82c7da532c4146f5408 Mon Sep 17 00:00:00 2001
From: Wen Gu <guwen@linux.alibaba.com>
Date: Tue, 19 Mar 2024 09:58:13 +0800
Subject: [PATCH] bpf_smc/smc-ebpf: add operations for strategy based on client
 IPv4 address

This patch includes:

- Patch #1:

bpf_smc: Add strategy based on IPv4 address information on the client side

1. destination address (that is the server address)

   Destination address is generally available. Check whether this address
   is allowed to use smc. If it is not explicitly set, further check whether
   0.0.0.0 should use smc. If the behavior of 0.0.0.0 is also not defined,
   it is allowed to use smc by default;

2. source address (that is the client address)

   If the kernel has set source address, check whether it is allowed to use
   smc. The check logic is the same as above. Otherwise, ignore this check;

- Patch #2:

smc-ebpf: add operations to configure and delete strategy based on client address.

- Patch #3:

smc-ebpf: add sub-command to check version.

- Patch #4:

smc-ebpf: use bpftool in current directory if not installed

- Patch #5:

smc-ebpf: refactor variables and functions to make code clear

Signed-off-by: Wen Gu <guwen@linux.alibaba.com>
---
 bpf_smc.c |  71 ++++++++-
 smc-ebpf  | 420 ++++++++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 430 insertions(+), 61 deletions(-)

diff --git a/bpf_smc.c b/bpf_smc.c
index 62f66ae..944194f 100644
--- a/bpf_smc.c
+++ b/bpf_smc.c
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0-only
-
 #include "vmlinux.h"
+#include <bpf/bpf_endian.h>
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_core_read.h>
 #include <bpf/bpf_tracing.h>
@@ -58,6 +58,25 @@ struct {
 	__type(value, struct smc_strategy);
 } smc_strategies SEC(".maps");
 
+struct smc_strat_ip_key {
+	__u32	prefix_len;	/* see struct bpf_lpm_trie_key */
+	__u8	ipv4[4];
+};
+
+struct smc_strat_ip_value {
+	__u8	mode;
+};
+
+/* maps for smc_strats_ip */
+struct {
+	__uint(type, BPF_MAP_TYPE_LPM_TRIE);
+	/* prefix_len & ip */
+	__uint(key_size, sizeof(struct smc_strat_ip_key));
+	__uint(value_size, sizeof(struct smc_strat_ip_value));
+	__uint(max_entries, 128);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+} smc_strats_ip SEC(".maps");
+
 struct smc_prediction {
 	/* count to allow smc */
 	__u32 credits;
@@ -109,11 +128,36 @@ static inline struct smc_strategy *smc_strategy_get(__u16 key)
 	return strategy;
 }
 
+static inline struct smc_strat_ip_value *smc_strategy_ip_get(__u32 ipv4)
+{
+	struct smc_strat_ip_key key = {
+		.prefix_len = 32,
+		.ipv4[3] = ipv4 & 0xFF,
+		.ipv4[2] = (ipv4 >> 8) & 0xFF,
+		.ipv4[1] = (ipv4 >> 16) & 0xFF,
+		.ipv4[0] = (ipv4 >> 24) & 0xFF,
+	};
+	struct smc_strat_ip_value *value;
+
+	value = bpf_map_lookup_elem(&smc_strats_ip, &key);
+	if (!value) {
+		/* check default dehavior */
+		key.ipv4[3] = 0;
+		key.ipv4[2] = 0;
+		key.ipv4[1] = 0;
+		key.ipv4[0] = 0;
+		value = bpf_map_lookup_elem(&smc_strats_ip, &key);
+	}
+	return value;
+}
+
 int SEC("struct_ops/bpf_smc_negotiate")
 BPF_PROG(bpf_smc_negotiate, struct sock *sk, struct sockaddr *addr)
 {
+	struct smc_strat_ip_value *ip_strat_value;
 	struct smc_prediction *smc_predictor;
 	struct smc_strategy *strategy;
+	__u32 ipv4_daddr, ipv4_saddr;
 	struct smc_sock *smc;
 	struct tcp_sock *tp;
 	__u64 now;
@@ -137,13 +181,32 @@ BPF_PROG(bpf_smc_negotiate, struct sock *sk, struct sockaddr *addr)
 		/* local port as key */
 		key = tp ? tp->inet_conn.icsk_inet.sk.__sk_common.skc_num : 0;
 
+#define DENYLIST_MODE	(0)
+#define AUTO_MODE	(1)
+#define ALLOWLIST_MODE	(2)
+	if (is_client) {
+		ipv4_daddr = addr ? bpf_ntohl(((struct sockaddr_in *) \
+					addr)->sin_addr.s_addr) : 0;
+		ip_strat_value = smc_strategy_ip_get(ipv4_daddr);
+		if (ip_strat_value &&
+		    ip_strat_value->mode == DENYLIST_MODE) {
+			return SK_DROP;
+		}
+		ipv4_saddr = tp ?
+			bpf_ntohl(tp->inet_conn.icsk_inet.sk.__sk_common.skc_rcv_saddr) : 0;
+		if (ipv4_saddr) {
+			ip_strat_value = smc_strategy_ip_get(ipv4_saddr);
+			if (ip_strat_value &&
+			    ip_strat_value->mode == DENYLIST_MODE) {
+				return SK_DROP;
+			}
+		}
+	}
+
 	strategy = smc_strategy_get(key);
 	if (!strategy)
 		return SK_DROP;
 
-#define DENYLIST_MODE	(0)
-#define AUTO_MODE		(1)
-#define ALLOWLIST_MODE	(2)
 	switch (strategy->mode) {
 	case AUTO_MODE:
 		if (!is_client)
diff --git a/smc-ebpf b/smc-ebpf
index f690137..60584ac 100755
--- a/smc-ebpf
+++ b/smc-ebpf
@@ -1,9 +1,12 @@
 #!/bin/bash
 
-BPFTOOL_EXEC="bpftool"
-TRAFFIC_CONTROL_NAME="anolis_smc"
-TRAFFIC_CONTROL_BINARY="anolis_smc"
-TRAFFIC_CONTROL_STRATEGIES_NAME="smc_strategies"
+VERSION="1.8.3"
+BPFT="bpftool"	# use system bpftool first if installed,
+                # otherwise, use bpftool in current directory.
+BPFPROG_NAME="anolis_smc"
+BPFPROG_FILE="anolis_smc"
+IP_POLICY="smc_strats_ip"
+PORT_POLICY="smc_strategies"
 
 default_config=(
     "0 1"
@@ -35,9 +38,11 @@ init_value() {
 }
 
 # 参数$1为要转换的数字，$2为字节数
+# Example: endian 0x18 4
+# output: 0x18 0x00 0x00 0x00
 endian() {
-    nums=$1
-    byte_size=$2
+    local nums=$1
+    local byte_size=$2
 
     if [ $(echo -n "1" | od -to2 | head -n1 | awk '{print $2}') = "0000000" ]; then
         for num in $(echo $nums | sed "s/,/ /g"); do
@@ -56,26 +61,216 @@ endian() {
     fi
 }
 
-smc_ebpf_usage() {
-    errMsg=$1
+# Convert IPv4 address to big endian __be32
+# Example: ipv4_aton 192.168.122.33
+# Output: 0x217AA8C0
+ipv4_aton() {
+	local IFS='.'
+	local ip=$1
+	local -a octets
+	local addr
+	local octet
+
+	# seperate IP to an array
+	read -ra octets <<< "$ip"
+
+	# check if IPv4 address includes 4 parts
+	if [ ${#octets[@]} -ne 4 ]; then
+		echo "error"
+		return
+	fi
+
+	# check if each part is between 0-255
+	for octet in "${octets[@]}"; do
+		if ! [[ "$octet" =~ ^[0-9]+$ ]]; then
+			echo "error"
+			return
+		fi
+		if [ "$octet" -lt 0 ] || [ "$octet" -gt 255 ]; then
+			echo "error"
+			return
+		fi
+	done
+
+	addr=$(( (${octets[3]} << 24) | (${octets[2]} << 16) | \
+		(${octets[1]} << 8) | ${octets[0]} ))
+	printf "0x%X\n" "$(( addr & 0xFFFFFFFF ))"
+}
+
+ipv4_verify_mask() {
+	local mask="$1"
+
+	if [ -z "$mask" ]; then
+		mask="32"	# if not set, use default '/32'
+	fi
+	if ! [[ "$mask" =~ ^[0-9]+$ ]]; then
+		echo "error"
+		return
+	fi
+	if [ "$mask" -lt 0 ] || [ "$mask" -gt 32 ]; then
+		echo "error"
+		return
+	fi
+	echo "$mask"
+}
+
+# Example usage: ipv4_add_map "192.168.122.33" "24" "02"
+ipv4_add_map() {
+	local ipv4="$1"
+	local ipv4_hex
+	local ipv4_mask="$2"
+	local mode="$3"
+
+	ipv4_hex=$(ipv4_aton "$ip")
+	if [[ ${ipv4_hex} == "error" ]]; then
+		echo "IP format is incorrect. Only IPv4 addresses are accepted"
+		exit 1
+	fi
+	ipv4_mask=$(ipv4_verify_mask "$ipv4_mask")
+	if [[ ${ipv4_mask} == "error" ]]; then
+		echo "IP mask is incorrect. Please use correct IPv4 mask"
+		exit 1
+	fi
+	# Example:
+	# bpftool map update name smc_starts_ip \
+	#	key 0x18 0x00 0x00 0x00 0xC0 0xA8 0x7A 0x21 value 0x02
+	${BPFT} map update name ${IP_POLICY} \
+		key $(endian $ipv4_mask 4) $(endian $ipv4_hex 4) \
+		value $(endian $mode 1)
+}
+
+# Example output of 'bpftool map dump name smc_strats_ip'
+#   # bpftool map dump name smc_strats_ip
+#   key: 10 00 00 00 c0 a7 03 00  value: 02
+#   key: 18 00 00 00 c0 a8 7a 21  value: 02
+#   key: 10 00 00 00 c0 a8 03 01  value: 02
+#   Found 3 elements
+#
+# We need to change it to
+#   key:     192.167.3.0/16       value:   "pass"
+#   key:     192.168.122.33/24    value:   "pass"
+#   key:     192.168.3.1/16       value:   "pass"
+#
+# Example usage: bpftool map dump name smc_strats_ip | ipv4_dump_map
+ipv4_dump_map() {
+	local format="%-8s %-20s %-8s %-10s\n"
+
+        while read -r line; do
+                # Check if the line contains the keys 'key' and 'value'
+                if echo "$line" | grep -q 'key:' && \
+                        echo "$line" | grep -q 'value:'; then
+                        # Extract the subnet prefix and IP address
+                        mask_hex=$(echo "$line" | awk '{
+                                for (i=1; i<=NF; i++)
+                                        if ($i == "key:")
+                                                print $(i+4)$(i+3)$(i+2)$(i+1)
+                                }')
+
+                        ip_hex=$(echo "$line" | awk '{
+                                for (i=1; i<=NF; i++)
+                                        if ($i == "key:")
+                                                print $(i+5)$(i+6)$(i+7)$(i+8)
+                                }')
+                        ip=$(printf "%d.%d.%d.%d/%d\n" \
+                                0x${ip_hex:0:2} 0x${ip_hex:2:2} \
+                                0x${ip_hex:4:2} 0x${ip_hex:6:2} \
+                                0x$mask_hex)
+
+                        # Extract and convert the value
+                        value_hex=$(echo "$line" | awk '{
+                                for (i=1; i<=NF; i++)
+                                        if ($i == "value:")
+                                                print $(i+1)
+                                }')
+                        case "$value_hex" in
+                                "00") value_str="denied" ;;
+                                "01") value_str="auto" ;;
+                                "02") value_str="pass" ;;
+                                *) value_str="unknown" ;;
+                        esac
+
+                        # Output the result
+                        printf "$format" \
+                                "key:" "${ip}" "value:" "\"${value_str}\""
+                fi
+        done
+}
+
+# Example usage: ipv4_delete_map "192.168.122.33" "24"
+ipv4_delete_map() {
+	local ipv4="$1"
+	local ipv4_hex
+	local ipv4_mask="$2"
+
+	ipv4_hex=$(ipv4_aton "$ip")
+	if [[ ${ipv4_hex} == "error" ]]; then
+		echo "IP format is incorrect. Only IPv4 addresses are accepted"
+		exit 1
+	fi
+	ipv4_mask=$(ipv4_verify_mask "$ipv4_mask")
+	if [[ ${ipv4_mask} == "error" ]]; then
+		echo "IP mask is incorrect. Please use correct IPv4 mask"
+		exit 1
+	fi
+
+	# Example:
+	# bpftool map delete name smc_strats_ip key \
+	#	0x18 0x00 0x00 0x00 0xc0 0xa8 0x7a 0x21
+	${BPFT} map delete name ${IP_POLICY} \
+		key $(endian $ipv4_mask 4) $(endian $ipv4_hex 4)
+}
+
+# Example output of 'bpftool map dump name smc_strats_ip'
+#   # bpftool map dump name smc_strats_ip
+#   key: 10 00 00 00 c0 a7 03 00  value: 02
+#   key: 18 00 00 00 c0 a8 7a 21  value: 02
+#   key: 10 00 00 00 c0 a8 03 01  value: 02
+#   Found 3 elements
+#
+# Example usage: bpftool map dump name smc_strats_ip | ipv4_delete_map_raw
+ipv4_delete_map_raw() {
+	local key_hex_array
+
+	while read -r line; do
+		if echo "$line" | grep -q 'key:'; then
+			# extract the next 8 bytes after 'key:'
+			key_hex_array=$(echo "$line" | awk '/key:/ {
+			for (i=2; i<=9; i++)
+				printf "0x%s ", $i
+				print ""
+			}')
+			# Example
+			# bpftool map delete name smc_strats_ip key \
+			#	0x18 0x00 0x00 0x00 0xc0 0xa8 0x7a 0x21
+			${BPFT} map delete \
+				name ${IP_POLICY} \
+				key $key_hex_array
+		fi
+	done
+}
+
+usage() {
+    local errMsg=$1
+
     if [ x"${errMsg}" != x ]; then
-        echo "Err: ${errMsg}"
+        echo "smc-ebpf: error: ${errMsg}"
     fi
 
     echo " Usage: smc-ebpf COMMAND"
     echo "        smc-ebpf policy               policy for connecting smc"
+    echo "        smc-ebpf version              display version information"
     echo " Examples: "
     echo "    smc-ebpf policy"
 
-    exit -1
+    exit 1
 }
 
 policy_usage() {
 
-    errMsg=$1
+    local errMsg=$1
 
     if [ x"${errMsg}" != x ]; then
-        echo "Err: ${errMsg}"
+        echo "smc-ebpf: error: ${errMsg}"
     fi
 
     echo " Usage: smc-ebpf policy COMMAND [OPTIONS]"
@@ -87,14 +282,18 @@ policy_usage() {
     echo "        smc-ebpf policy clear             clear all policy config"
     echo "        smc-ebpf policy dump              display all policy config"
     echo "        smc-ebpf policy config [OPTIONS]  config policy"
+    echo "        smc-ebpf policy delete [OPTIONS]  delete policy"
+    echo "    --ip [IPv4]                           target IPv4 address"
     echo "    --port                                target port"
     echo "    --mode [auto|disable|enable]          target mode"
     echo " Examples: "
     echo "    smc-ebpf policy load"
     echo "    #disable port 80 to use smc"
     echo "    smc-ebpf policy config --port 80 --mode disable "
+    echo "    #delete ip 192.168.0.0/24 policy"
+    echo "    smc-ebpf policy delete --ip 192.168.0.0 --mask 24 "
 
-    exit -1
+    exit 1
 }
 
 policy_help() {
@@ -103,10 +302,11 @@ policy_help() {
 
 policy_config() {
 
-    opts="port:,mode:,auto,enable,disable"
+    local opts="port:,ip:,mask:,mode:,auto,enable,disable"
+    local port ip mask mode
 
     # parse config
-    parsed_args=$(getopt --long ${opts} -- "$@")
+    local parsed_args=$(getopt --long ${opts} -- "$@")
     eval set -- "$parsed_args"
     while [ -n "$1" ]; do
         case "$1" in
@@ -114,16 +314,29 @@ policy_config() {
             port=$2
             shift 2
             ;;
+	--ip)
+	    ip=$2
+	    shift 2
+	    ;;
+	--mask)
+	    mask=$2
+	    shift 2
+	    ;;
         --mode)
-            if [ x"$2" == x"auto" ]; then
-                mode=1
-            elif [ x"$2" == x"enable" ]; then
-                mode=2
-            elif [ x"$2" == x"disable" ]; then
-                mode=0
-            else
-                policy_usage "unrecognized mode: $2"
-            fi
+            case "$2" in
+                auto)
+                    mode=1
+                    ;;
+                enable)
+                    mode=2
+                    ;;
+                disable)
+                    mode=0
+                    ;;
+                *)
+                    policy_usage "unrecognized mode: $2"
+                    ;;
+            esac
             shift 2
             ;;
         --auto)
@@ -148,41 +361,107 @@ policy_config() {
         esac
     done
 
+    if [ -n "$ip" ] && [ -n "$port" ]; then
+	policy_usage "Error: ip and port cannot be set at the same time."
+	exit 1
+    fi
+
     # update config
-    init_value
-    ${BPFTOOL_EXEC} map update name ${TRAFFIC_CONTROL_STRATEGIES_NAME} key $(endian $port 2) value \
-        $(endian $mode 1) 0x00 \
-        $(endian $rtt_threshold 2) \
-        $(endian $smc_productivity 2) \
-        $(endian $tcp_productivity 2) \
-        $(endian $max_credits 4) \
-        $(endian $min_credits 4) \
-        $(endian $initial_credits 4) \
-        $(endian $max_pacing_burst 4) \
-        0x00 0x00 0x00 0x00 \
-        $(endian $pacing_delta 8)
+    if [ -n "$ip" ]; then
+	    ipv4_add_map "$ip" "$mask" "$mode"
+    fi
+
+    if [ -n "$port" ]; then
+	    init_value
+	    ${BPFT} map update name ${PORT_POLICY} \
+		key $(endian $port 2) \
+		value \
+		$(endian $mode 1) 0x00 \
+		$(endian $rtt_threshold 2) \
+		$(endian $smc_productivity 2) \
+		$(endian $tcp_productivity 2) \
+		$(endian $max_credits 4) \
+		$(endian $min_credits 4) \
+		$(endian $initial_credits 4) \
+		$(endian $max_pacing_burst 4) \
+		0x00 0x00 0x00 0x00 \
+		$(endian $pacing_delta 8)
+    fi
+}
+
+policy_delete() {
+
+    local opts="port:,ip:,mask:"
+    local port ip mask
+
+    # parse config
+    local parsed_args=$(getopt --long ${opts} -- "$@")
+    eval set -- "$parsed_args"
+    while [ -n "$1" ]; do
+        case "$1" in
+        --port)
+            port=$2
+            shift 2
+            ;;
+	--ip)
+	    ip=$2
+	    shift 2
+	    ;;
+	--mask)
+	    mask=$2
+	    shift 2
+	    ;;
+        --)
+            shift
+            break
+            ;;
+        *)
+            policy_usage "unrecognized OPTION: $1"
+            ;;
+        esac
+    done
+
+    if [ -n "$ip" ] && [ -n "$port" ]; then
+        policy_usage "Error: IP and port cannot be set at the same time."
+        exit 1
+    fi
+
+    # update config
+    if [ -n "$ip" ]; then
+	    ipv4_delete_map "$ip" "$mask" "$mode"
+    fi
+
+    if [ -n "$port" ]; then
+	    ${BPFT} map delete name ${PORT_POLICY} key $(endian $port 2)
+    fi
 }
 
 # clear all config
 policy_clear() {
-    local keys=$(${BPFTOOL_EXEC} map dump name ${TRAFFIC_CONTROL_STRATEGIES_NAME} | grep -oP '(?<= "key": )\d+')
-    for key in ${keys[@]}; do
-        ${BPFTOOL_EXEC} map delete name ${TRAFFIC_CONTROL_STRATEGIES_NAME} key $(endian $key 2)
+    # clear port policy
+    local keys=$(${BPFT} map dump name ${PORT_POLICY} \
+		| grep -oP '(?<= "key": )\d+')
+    for key in "${keys[@]}"; do
+        ${BPFT} map delete name ${PORT_POLICY} key $(endian $key 2)
     done
+
+    # clear IPv4 policy
+    ${BPFT} map dump name ${IP_POLICY} | ipv4_delete_map_raw
 }
 
-# destroy traffic control
+# unload ebpf program
 policy_unload() {
     policy_clear
-    # unregister
-    ${BPFTOOL_EXEC} struct_ops unregister name ${TRAFFIC_CONTROL_NAME}
+    ${BPFT} struct_ops unregister name ${BPFPROG_NAME}
 }
 
 policy_init() {
     for def in "${default_config[@]}"; do
         eval set -- ${def}
         init_value
-        ${BPFTOOL_EXEC} map update name ${TRAFFIC_CONTROL_STRATEGIES_NAME} key $(endian $1 2) value \
+        ${BPFT} map update name ${PORT_POLICY} \
+	    key $(endian $1 2) \
+	    value \
             $(endian $2 1) 0x00 \
             $(endian $rtt_threshold 2) \
             $(endian $smc_productivity 2) \
@@ -194,18 +473,23 @@ policy_init() {
             0x00 0x00 0x00 0x00 \
             $(endian $pacing_delta 8)
     done
+    # ip strategies won't be involved.
 }
 
 policy_dump() {
-    ${BPFTOOL_EXEC} map dump name ${TRAFFIC_CONTROL_STRATEGIES_NAME}
+    # dump IPv4 policy
+    ${BPFT} map dump name ${IP_POLICY} \
+	| ipv4_dump_map
+    # dump port policy
+    ${BPFT} map dump name ${PORT_POLICY}
 }
 
-# load a traffic control ebpf-file
+# load ebpf program
 policy_load() {
 
     local init=0
+    local opts="init"
 
-    opts="init"
     # parse config
     parsed_args=$(getopt --long ${opts} -- "$@")
     eval set -- "$parsed_args"
@@ -225,7 +509,7 @@ policy_load() {
     done
 
     #check if ebpf prog was already loaded
-    ${BPFTOOL_EXEC} struct_ops | grep ${TRAFFIC_CONTROL_NAME} >/dev/null 2>&1
+    ${BPFT} struct_ops | grep ${BPFPROG_NAME} >/dev/null 2>&1
     if [ x"$?" == x'0' ]; then
         if [ $init -eq 1 ]; then
             policy_init
@@ -234,7 +518,7 @@ policy_load() {
     fi
 
     local DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
-    ${BPFTOOL_EXEC} struct_ops register $DIR/${TRAFFIC_CONTROL_BINARY}
+    ${BPFT} struct_ops register $DIR/${BPFPROG_FILE}
 
     if [ $init -eq 1 ]; then
 	policy_init
@@ -248,25 +532,47 @@ policy_op() {
         policy_usage "required root privileged"
     fi
 
+    # adjust bpftool path
+    if command -v bpftool >/dev/null 2>&1; then
+        BPFT=$(command -v bpftool)
+    else
+        if [ -x "./bpftool" ]; then
+	    BPFT="./bpftool"
+        else
+	    policy_usage "bpftool is not installed."
+        fi
+    fi
+
+
     case $op in
-    load | unload | stop | init | clear | config | help | dump) ;;
+    load | unload | stop | init | clear | config | delete | help | dump) ;;
+    "")
+	policy_usage
+	;;
     *)
         policy_usage "unrecognized COMMAND: ${op}"
         ;;
     esac
-    eval policy_${op} "${op} $@"
+    eval policy_${op} "${op}" "$@"
 }
 
+# main
 obj="$1"
 shift
-op="$1"
+op="$1"	# global
 shift
 
 case $obj in
-policy)
-    policy_op "$@"
-    ;;
-*)
-    smc_ebpf_usage "unrecognized ${obj}"
-    ;;
+    policy)
+        policy_op "$@"
+        ;;
+    version)
+        echo "smc-ebpf utility, smc-tools-$VERSION"
+        ;;
+    "")
+        usage
+        ;;
+    *)
+        usage "unrecognized: ${obj}"
+        ;;
 esac
-- 
2.43.5