From d308722d4865fef0339751a1663652c586137054 Mon Sep 17 00:00:00 2001
From: r12f <r12f.code@gmail.com>
Date: Tue, 1 Oct 2024 20:12:32 -0700
Subject: [PATCH] Remove tunnel common setting from next hop.

---
 dash-pipeline/bmv2/stages/tunnel_stage.p4 | 67 +++++++++++++----------
 1 file changed, 38 insertions(+), 29 deletions(-)

diff --git a/dash-pipeline/bmv2/stages/tunnel_stage.p4 b/dash-pipeline/bmv2/stages/tunnel_stage.p4
index e1b7bfe9c..23578e63a 100644
--- a/dash-pipeline/bmv2/stages/tunnel_stage.p4
+++ b/dash-pipeline/bmv2/stages/tunnel_stage.p4
@@ -6,19 +6,21 @@ control tunnel_stage(
     inout metadata_t meta)
 {
     action set_tunnel_attrs(
-        @SaiVal[type="sai_ip_address_t"]
-        IPv4Address dip,
         @SaiVal[type="sai_dash_encapsulation_t", default_value="SAI_DASH_ENCAPSULATION_VXLAN"]
         dash_encapsulation_t dash_encapsulation,
-        bit<24> tunnel_key)
+
+        bit<24> tunnel_key,
+
+        @SaiVal[type="sai_ip_address_t"]
+        IPv4Address dip,
+
+        @SaiVal[type="sai_ip_address_t"]
+        IPv4Address sip)
     {
-        push_action_static_encap(hdr = hdr,
-                                 meta = meta,
-                                 encap = dash_encapsulation,
-                                 vni = tunnel_key,
-                                 underlay_sip = hdr.u0_ipv4.src_addr,
-                                 underlay_dip = dip,
-                                 overlay_dmac = hdr.u0_ethernet.dst_addr);
+        meta.tunnel_data.dash_encapsulation = dash_encapsulation;
+        meta.tunnel_data.vni = tunnel_key;
+        meta.tunnel_data.underlay_sip = sip == 0 ? hdr.u0_ipv4.src_addr : sip;
+        meta.tunnel_data.underlay_dip = dip;
     }
 
     @SaiTable[name = "dash_tunnel", api = "dash_tunnel", order = 0, isobject="true"]
@@ -57,7 +59,7 @@ control tunnel_stage(
         @SaiVal[type="sai_object_id_t"] bit<16> dash_tunnel_next_hop_id)
     {
         // dash_tunnel_id in tunnel member must match the metadata
-        assert(meta.dash_tunnel_id == dash_tunnel_id);
+        REQUIRES(meta.dash_tunnel_id == dash_tunnel_id);
 
         meta.dash_tunnel_next_hop_id = dash_tunnel_next_hop_id;
     }
@@ -75,22 +77,13 @@ control tunnel_stage(
 
     action set_tunnel_next_hop_attrs(
         @SaiVal[type="sai_ip_address_t"]
-        IPv4Address dip,
-        @SaiVal[type="sai_dash_encapsulation_t", default_value="SAI_DASH_ENCAPSULATION_VXLAN"]
-        dash_encapsulation_t dash_encapsulation,
-        bit<24> tunnel_key)
+        IPv4Address dip)
     {
-        push_action_static_encap(hdr = hdr,
-                                 meta = meta,
-                                 encap = dash_encapsulation,
-                                 vni = tunnel_key,
-                                 underlay_sip = hdr.u0_ipv4.src_addr,
-                                 underlay_dip = dip,
-                                 overlay_dmac = hdr.u0_ethernet.dst_addr);
+        meta.tunnel_data.underlay_dip = dip == 0 ? meta.tunnel_data.underlay_dip : dip;
     }
 
     @SaiTable[name = "dash_tunnel_next_hop", api = "dash_tunnel", order = 2, isobject="true"]
-    table tunnel_next_hop{
+    table tunnel_next_hop {
         key = {
             meta.dash_tunnel_next_hop_id : exact @SaiVal[type="sai_object_id_t"];
         }
@@ -101,15 +94,31 @@ control tunnel_stage(
     }
 
     apply {
+        if (meta.dash_tunnel_id == 0) {
+            return;
+        }
+
         tunnel.apply();
 
-        // TODO: Calculate based on packet and tunnel member size.
-        // Currently, we will have to use 0 here, because we don't know how many tunnel members we will have.
-        meta.dash_tunnel_member_index = 0;
-        tunnel_member_select.apply();
+        if (meta.tunnel_data.underlay_dip == 0) {
+            // If underlay_dip is not set by the tunnel, we are using the tunnel as ECMP group.
 
-        tunnel_member.apply();
-        tunnel_next_hop.apply();
+            // TODO: Calculate based on packet and tunnel member size.
+            // Currently, we will have to use 0 here, because we don't know how many tunnel members we will have.
+            meta.dash_tunnel_member_index = 0;
+            tunnel_member_select.apply();
+
+            tunnel_member.apply();
+            tunnel_next_hop.apply();
+        }
+
+        push_action_static_encap(hdr = hdr,
+                                 meta = meta,
+                                 encap = meta.tunnel_data.dash_encapsulation,
+                                 vni = meta.tunnel_data.vni,
+                                 underlay_sip = meta.tunnel_data.underlay_sip,
+                                 underlay_dip = meta.tunnel_data.underlay_dip,
+                                 overlay_dmac = hdr.u0_ethernet.dst_addr);
     }
 }