From 8d81db65bc05714f98e61c9d2ee4434eb2a536a4 Mon Sep 17 00:00:00 2001
From: gpunathilell <gpunathilell@nvidia.com>
Date: Fri, 2 Aug 2024 22:50:25 +0000
Subject: [PATCH 1/5] Added inbound traffic capability for script

---
 files/scripts/sonic-dpu-mgmt-traffic.sh | 228 ++++++++++++++++++++++--
 1 file changed, 211 insertions(+), 17 deletions(-)

diff --git a/files/scripts/sonic-dpu-mgmt-traffic.sh b/files/scripts/sonic-dpu-mgmt-traffic.sh
index 71e6ed29b32..9c6f7b06952 100755
--- a/files/scripts/sonic-dpu-mgmt-traffic.sh
+++ b/files/scripts/sonic-dpu-mgmt-traffic.sh
@@ -4,12 +4,22 @@
 command_name=$0
 
 usage(){
-    echo "Syntax: $command_name -e|--enable -d|--disable"
+    echo "Syntax: $command_name inbound/outbound -e|--enable -d|--disable [--dpus,--ports,--nofwctrl]"
     echo "Arguments:"
-    echo "-e			Enable dpu management traffic forwarding"
-    echo "-d			Disable dpu management traffic forwarding"
+    echo "inbound		Control DPU Inbound traffic forwarding"
+    echo "outbound		Control DPU Outbound traffic forwarding"
+    echo "-e			Enable dpu management traffic forwarding in a specific direction"
+    echo "-d			Disable dpu management traffic forwarding in a specific direction"
+    echo "--dpus		Selection of dpus for which the inbound traffic has to be controlled (all can be specified)"
+    echo "--ports		Selection of ports for which the inbound traffic has to be controlled"
+    echo "--nofwctrl	Disable changing the general ipv4 forwarding (Could be useful if inbound is enabled and outbound is disabled)"
 }
 
+dpu_l=()
+declare -A midplane_dict
+declare -A midplane_ip_dict
+fw_change="enable"
+
 add_rem_valid_iptable(){
     local op=$1
     local table=$2
@@ -39,20 +49,135 @@ control_forwarding(){
     if [ "$op" = "enable" ]; then
         value=1
     fi
-    echo $value > /proc/sys/net/ipv4/ip_forward
-    echo $value > /proc/sys/net/ipv4/conf/eth0/forwarding
+    if [ "$fw_change" = "enable" ]; then
+        echo $value > /proc/sys/net/ipv4/ip_forward
+        echo $value > /proc/sys/net/ipv4/conf/eth0/forwarding
+    fi
+}
+
+validate_dpus(){
+    local provided_list=("$@")
+    for item1 in "${provided_list[@]}"; do
+        local found=0
+        for item2 in "${dpu_l[@]}"; do
+            if [[ "$item1" = "$item2" ]]; then
+                found=1
+                break
+            fi
+        done
+        if [[ $found -eq 0 ]]; then
+            echo "$item1 is not detected! Please provide proper dpu names list!"
+            exit 1
+        fi
+    done
 }
 
-ctrl_dpu_forwarding(){
+general_validation(){
+    if [ -z "$direction" ]; then
+        echo "Please provide the direction argument (inbound or outboud)"
+        usage
+        exit 1
+    fi
+
+    if [ -z "$operation" ]; then
+        echo "Please provide the operation option (-e or -d)"
+        usage
+        exit 1
+    fi
+}
+
+
+inbound_validation(){
+    #DPU Validation
+    while IFS= read -r line; do
+        dpu_name=$( echo "$line" | sed 's/.*|//;s/"//g')
+        dpu_l+=("$dpu_name")
+    done < <(redis-cli -n 4 keys DPUS*)
+    len1=${#dpu_name[@]}
+    if [ "$len1" -eq 0 ]; then
+        echo "No dpus detected on device!"
+        exit 1
+    fi
+    sorted_dpu_l=($(for item in "${dpu_l[@]}"; do
+        echo "$item"
+    done | sort))
+    if [ -z "$arg_dpu_names" ]; then
+            echo "No DPUs provided!"
+            usage
+            exit 1
+        else
+            if [ "$arg_dpu_names" = "all" ]; then
+                sel_dpu_names=("${sorted_dpu_l[@]}")
+                echo "${#sorted_dpu_l[@]} DPUs detected:"
+                echo "${sorted_dpu_l[@]}"
+            else
+                IFS=',' read -ra sel_dpu_names <<< "$arg_dpu_names"
+                validate_dpus ${sel_dpu_names[@]}
+            fi
+    fi
+    #Port validation
+    IFS=',' read -ra provided_ports <<< "$arg_port_list"
+    len1=${#sel_dpu_names[@]}
+    len2=${#provided_ports[@]}
+    if [ "$len1" -ne "$len2" ]; then
+        echo "Length of ${sel_dpu_names[@]} does not match provided port length ${provided_ports[@]}"
+        usage
+        exit 1
+    fi
+    for dpu in "${sel_dpu_names[@]}"; do
+        midplane_int_name=$(redis-cli -n 4 hget "DPUS|$dpu" "midplane_interface")
+        if [ -z "$midplane_int_name" ]; then
+            echo "Cannot obtain midplane interface for $dpu"
+            exit 1
+        fi
+        midplane_dict["$dpu"]="$midplane_int_name"
+    done
+
+    for key in "${!midplane_dict[@]}"; do
+        echo "\"$key\":\"${midplane_dict[$key]}\""
+    done
+
+    for dpu in "${!midplane_dict[@]}"; do
+        midplane_ip=$(redis-cli -n 4 hget "DHCP_SERVER_IPV4_PORT|$midplane_iface|${midplane_dict[$dpu]}" "ips@")
+        if [ -z "$midplane_ip" ]; then
+            echo "Cannot obtain midplane ip for $dpu"
+            exit 1
+        fi
+        midplane_ip_dict["$dpu"]="$midplane_ip"
+    done
+}
+
+# Outbound Traffice forwarding control function
+ctrl_dpu_ob_forwarding(){
     local op=$1
     control_forwarding $op
     add_rem_valid_iptable $op nat POSTROUTING -o ${mgmt_iface}  -j MASQUERADE
     add_rem_valid_iptable $op filter FORWARD -i ${mgmt_iface} -o ${midplane_iface} -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
     add_rem_valid_iptable $op filter FORWARD -i ${midplane_iface} -o ${mgmt_iface} -j ACCEPT
     if [ "$op" = "enable" ]; then
-        echo "Enabled DPU management traffic Forwarding"
+        echo "Enabled DPU management outbound traffic Forwarding"
     else
-        echo "Disabled DPU management traffic Forwarding"
+        echo "Disabled DPU management outbound traffic Forwarding"
+    fi
+}
+
+ctrl_dpu_ib_forwarding(){
+    local op=$1
+    local dest_port=22
+    control_forwarding $op
+    add_rem_valid_iptable $op filter FORWARD -i ${midplane_iface} -o ${mgmt_iface} -p tcp -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
+    add_rem_valid_iptable $op filter FORWARD -i ${mgmt_iface} -o ${midplane_iface} -p tcp --dport $dest_port -m conntrack --ctstate NEW,RELATED,ESTABLISHED -j ACCEPT
+    for index in ${!sel_dpu_names[@]}; do
+        dpu_name="${sel_dpu_names[$index]}"
+        dpu_midplane_ip="${midplane_ip_dict[$dpu_name]}"
+        switch_port="${provided_ports[$index]}"
+        add_rem_valid_iptable $op nat POSTROUTING -p tcp -d $dpu_midplane_ip --dport $dest_port -j MASQUERADE
+        add_rem_valid_iptable $op nat PREROUTING -i ${mgmt_iface}  -p tcp --dport $switch_port -j DNAT --to-destination $dpu_midplane_ip:$dest_port
+    done
+    if [ "$op" = "enable" ]; then
+        echo "Enabled DPU management inbound traffic Forwarding"
+    else
+        echo "Disabled DPU management inbound traffic Forwarding"
     fi
 }
 
@@ -70,16 +195,85 @@ if ! ifconfig "$midplane_iface" > /dev/null 2>&1; then
     exit 1
 fi
 
+operation=""
+direction=""
+
+invalid_arg(){
+    echo "Invalid arguments $1"
+    usage
+    exit 1
+}
+
+
 case $1 in
-    -e|--enable)
-        ctrl_dpu_forwarding enable
-        ;;
-    -d|--disable)
-        ctrl_dpu_forwarding disable
-        ;;
+    inbound)
+        direction="inbound"
+        shift
+        while [ "$1" != "--" ]  && [ -n "$1" ]; do
+            case $1 in
+                -e|--enable)
+                    operation="enable"
+                ;;
+                -d|--disable)
+                    operation="disable"
+                ;;
+                --dpus)
+                    shift;
+                    arg_dpu_names=$1
+                ;;
+                --ports)
+                    shift;
+                    arg_port_list=$1
+                ;;
+                --ports)
+                    shift;
+                    arg_port_list=$1
+                ;;
+                --nofwctrl)
+                    fw_change="disable"
+                ;;
+                *)
+                    invalid_arg $1
+                ;;
+            esac
+        shift
+        done
+    ;;
+    outbound)
+        direction="outbound"
+        shift
+        while [ "$1" != "--" ]  && [ -n "$1" ]; do
+            case $1 in
+                -e|--enable)
+                    operation="enable"
+                ;;
+                -d|--disable)
+                    operation="disable"
+                ;;
+                --nofwctrl)
+                    fw_change="disable"
+                ;;
+                *)
+                    invalid_arg $1
+                ;;
+            esac
+        shift
+        done
+    ;;
     *)
-        echo "Incorrect Usage!"
-        usage
-        exit 1
+        invalid_arg $1
+    ;;
+esac
+
+general_validation
+
+
+case $direction in
+    outbound)
+        ctrl_dpu_ob_forwarding $operation
+        ;;
+    inbound)
+        inbound_validation
+        ctrl_dpu_ib_forwarding $operation
         ;;
 esac

From 0949551eebb0d0ee31924c30fd5571ad444fd8dc Mon Sep 17 00:00:00 2001
From: gpunathilell <gpunathilell@nvidia.com>
Date: Mon, 28 Oct 2024 10:03:52 +0000
Subject: [PATCH 2/5] Removed connection based filtering

---
 files/scripts/sonic-dpu-mgmt-traffic.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/files/scripts/sonic-dpu-mgmt-traffic.sh b/files/scripts/sonic-dpu-mgmt-traffic.sh
index 9c6f7b06952..a926ae5c60b 100755
--- a/files/scripts/sonic-dpu-mgmt-traffic.sh
+++ b/files/scripts/sonic-dpu-mgmt-traffic.sh
@@ -166,7 +166,7 @@ ctrl_dpu_ib_forwarding(){
     local dest_port=22
     control_forwarding $op
     add_rem_valid_iptable $op filter FORWARD -i ${midplane_iface} -o ${mgmt_iface} -p tcp -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
-    add_rem_valid_iptable $op filter FORWARD -i ${mgmt_iface} -o ${midplane_iface} -p tcp --dport $dest_port -m conntrack --ctstate NEW,RELATED,ESTABLISHED -j ACCEPT
+    add_rem_valid_iptable $op filter FORWARD -i ${mgmt_iface} -o ${midplane_iface} -p tcp --dport $dest_port -j ACCEPT
     for index in ${!sel_dpu_names[@]}; do
         dpu_name="${sel_dpu_names[$index]}"
         dpu_midplane_ip="${midplane_ip_dict[$dpu_name]}"

From d0ca225c5e7bc8e189dfa4a42f1dab1378007f11 Mon Sep 17 00:00:00 2001
From: gpunathilell <gpunathilell@nvidia.com>
Date: Mon, 28 Oct 2024 10:09:18 +0000
Subject: [PATCH 3/5] Removed additional option

---
 files/scripts/sonic-dpu-mgmt-traffic.sh | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/files/scripts/sonic-dpu-mgmt-traffic.sh b/files/scripts/sonic-dpu-mgmt-traffic.sh
index a926ae5c60b..63699b12e5a 100755
--- a/files/scripts/sonic-dpu-mgmt-traffic.sh
+++ b/files/scripts/sonic-dpu-mgmt-traffic.sh
@@ -225,10 +225,6 @@ case $1 in
                     shift;
                     arg_port_list=$1
                 ;;
-                --ports)
-                    shift;
-                    arg_port_list=$1
-                ;;
                 --nofwctrl)
                     fw_change="disable"
                 ;;

From 3ec3bdde597fc09e8f5db3ceaf3c2b20cb08c380 Mon Sep 17 00:00:00 2001
From: gpunathilell <gpunathilell@nvidia.com>
Date: Fri, 15 Nov 2024 17:27:49 +0000
Subject: [PATCH 4/5] Added inbound port validation

---
 files/scripts/sonic-dpu-mgmt-traffic.sh | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/files/scripts/sonic-dpu-mgmt-traffic.sh b/files/scripts/sonic-dpu-mgmt-traffic.sh
index 63699b12e5a..c588f029ea3 100755
--- a/files/scripts/sonic-dpu-mgmt-traffic.sh
+++ b/files/scripts/sonic-dpu-mgmt-traffic.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/bin/bash -x
 #Script to control the DPU management traffic forwarding through the SmartSwitch
 
 command_name=$0
@@ -86,6 +86,22 @@ general_validation(){
     fi
 }
 
+port_use_validation(){
+	local port_l=("$@")
+	for port in "${port_l[@]}"; do
+		if (( port >= 0 && port <= 1023 )); then
+			echo "Provided port $port in range 0-1023, Please execute with a different port"
+			exit 1
+		fi
+		if netstat -tuln | awk '{print $4}' | grep -q ":$port\$"; then
+			echo "Provided port $port is in use by another process, Please execute with a different port"
+			exit 1
+		fi
+	done
+}
+
+
+
 
 inbound_validation(){
     #DPU Validation
@@ -124,6 +140,7 @@ inbound_validation(){
         usage
         exit 1
     fi
+    port_use_validation ${provided_ports[@]}
     for dpu in "${sel_dpu_names[@]}"; do
         midplane_int_name=$(redis-cli -n 4 hget "DPUS|$dpu" "midplane_interface")
         if [ -z "$midplane_int_name" ]; then
@@ -133,10 +150,6 @@ inbound_validation(){
         midplane_dict["$dpu"]="$midplane_int_name"
     done
 
-    for key in "${!midplane_dict[@]}"; do
-        echo "\"$key\":\"${midplane_dict[$key]}\""
-    done
-
     for dpu in "${!midplane_dict[@]}"; do
         midplane_ip=$(redis-cli -n 4 hget "DHCP_SERVER_IPV4_PORT|$midplane_iface|${midplane_dict[$dpu]}" "ips@")
         if [ -z "$midplane_ip" ]; then

From b56bb49ce76f24f4f37247cd8f82f00a52cf55b2 Mon Sep 17 00:00:00 2001
From: gpunathilell <gpunathilell@nvidia.com>
Date: Fri, 15 Nov 2024 17:31:56 +0000
Subject: [PATCH 5/5] Removed debugging mode

---
 files/scripts/sonic-dpu-mgmt-traffic.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/files/scripts/sonic-dpu-mgmt-traffic.sh b/files/scripts/sonic-dpu-mgmt-traffic.sh
index c588f029ea3..014d722537f 100755
--- a/files/scripts/sonic-dpu-mgmt-traffic.sh
+++ b/files/scripts/sonic-dpu-mgmt-traffic.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -x
+#!/bin/bash
 #Script to control the DPU management traffic forwarding through the SmartSwitch
 
 command_name=$0