From 03a10466e1876886ec91c70caaa071630c9916b3 Mon Sep 17 00:00:00 2001
From: Marian Pritsak <marianp@mellanox.com>
Date: Tue, 21 Jan 2025 07:25:51 +0200
Subject: [PATCH] Override VM MAC when in DST address ENI lookup

Signed-off-by: Marian Pritsak <marianp@mellanox.com>
---
 dash-pipeline/bmv2/dash_conntrack.p4    | 5 +++--
 dash-pipeline/bmv2/dash_inbound.p4      | 7 +++++++
 dash-pipeline/bmv2/dash_metadata.p4     | 1 +
 dash-pipeline/bmv2/dash_outbound.p4     | 5 +++++
 dash-pipeline/bmv2/stages/eni_lookup.p4 | 3 ++-
 5 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/dash-pipeline/bmv2/dash_conntrack.p4 b/dash-pipeline/bmv2/dash_conntrack.p4
index 1b6195920..9cc191937 100644
--- a/dash-pipeline/bmv2/dash_conntrack.p4
+++ b/dash-pipeline/bmv2/dash_conntrack.p4
@@ -39,7 +39,7 @@ control ConntrackIn(inout headers_t hdr,
                 inout metadata_t meta)
 {
 
-  action conntrackIn_allow (IPv4Address original_overlay_sip, IPv4Address original_overlay_dip) {
+  action conntrackIn_allow (IPv4Address original_overlay_sip, IPv4Address original_overlay_dip, EthernetAddress vm_nic_addr) {
   /* Invalidate entry based on TCP flags */
           // If FIN is 1 (0b000001), or if RST is 1 (0b000100):
           if ((hdr.customer_tcp.flags & 0b000101 /* FIN/RST */) != 0) {
@@ -51,6 +51,7 @@ control ConntrackIn(inout headers_t hdr,
           meta.overlay_data.is_ipv6 = 0;
           meta.overlay_data.sip = (IPv4ORv6Address)original_overlay_sip;
           meta.overlay_data.dip = (IPv4ORv6Address)original_overlay_dip;
+          meta.vm_nic_addr = vm_nic_addr;
   }
 
   action conntrackIn_miss() {
@@ -60,7 +61,7 @@ control ConntrackIn(inout headers_t hdr,
                 if (meta.routing_actions & dash_routing_actions_t.NAT46 != 0) {
                     // New PNA Extern
                     add_entry("conntrackIn_allow",
-                                {(IPv4Address)meta.src_ip_addr, (IPv4Address)meta.dst_ip_addr},
+                                {(IPv4Address)meta.src_ip_addr, (IPv4Address)meta.dst_ip_addr, meta.vm_nic_addr},
                                 EXPIRE_TIME_PROFILE_LONG);
                 }
                 //adding failure to be eventually handled
diff --git a/dash-pipeline/bmv2/dash_inbound.p4 b/dash-pipeline/bmv2/dash_inbound.p4
index f96c3029e..316df826d 100644
--- a/dash-pipeline/bmv2/dash_inbound.p4
+++ b/dash-pipeline/bmv2/dash_inbound.p4
@@ -46,7 +46,14 @@ control inbound(inout headers_t hdr,
                      meta.u0_encap_data.underlay_sip,
                      dash_encapsulation_t.VXLAN,
                      meta.u0_encap_data.vni);
+
+        // Fix up overlay DST MAC to be VM MAC and SRC MAC to ENI MAC
+        if (meta.eni_mac_type == dash_eni_mac_type_t.DST_MAC) {
+            hdr.customer_ethernet.dst_addr = meta.vm_nic_addr;
+            hdr.customer_ethernet.src_addr = meta.eni_addr;
+        }
     }
+
 }
 
 #endif /* _SIRIUS_INBOUND_P4_ */
diff --git a/dash-pipeline/bmv2/dash_metadata.p4 b/dash-pipeline/bmv2/dash_metadata.p4
index bcbbe8ba6..a8b439663 100644
--- a/dash-pipeline/bmv2/dash_metadata.p4
+++ b/dash-pipeline/bmv2/dash_metadata.p4
@@ -176,6 +176,7 @@ struct metadata_t {
     encap_data_t rx_encap;
 #endif // TARGET_DPDK_PNA
     EthernetAddress eni_addr;
+    EthernetAddress vm_nic_addr;
     bit<16> vnet_id;
     bit<16> dst_vnet_id;
     bit<16> eni_id;
diff --git a/dash-pipeline/bmv2/dash_outbound.p4 b/dash-pipeline/bmv2/dash_outbound.p4
index 8440b7a84..e4903ff7e 100644
--- a/dash-pipeline/bmv2/dash_outbound.p4
+++ b/dash-pipeline/bmv2/dash_outbound.p4
@@ -40,6 +40,11 @@ control outbound(inout headers_t hdr,
         outbound_mapping_stage.apply(hdr, meta);
 
         outbound_pre_routing_action_apply_stage.apply(hdr, meta);
+
+        // Fix up overlay SRC MAC to be VM MAC
+        if (meta.eni_mac_type == dash_eni_mac_type_t.DST_MAC) {
+            hdr.customer_ethernet.src_addr = meta.eni_addr;
+        }
     }
 }
 
diff --git a/dash-pipeline/bmv2/stages/eni_lookup.p4 b/dash-pipeline/bmv2/stages/eni_lookup.p4
index 5c9352013..4b950c628 100644
--- a/dash-pipeline/bmv2/stages/eni_lookup.p4
+++ b/dash-pipeline/bmv2/stages/eni_lookup.p4
@@ -32,6 +32,7 @@ control eni_lookup_stage(
             meta.eni_addr = hdr.customer_ethernet.src_addr;
         } else {
             meta.eni_addr = hdr.customer_ethernet.dst_addr;
+            meta.vm_nic_addr = hdr.customer_ethernet.src_addr;
         }
                                           
         if (!eni_ether_address_map.apply().hit) {
@@ -44,4 +45,4 @@ control eni_lookup_stage(
     }
 }
 
-#endif /* _DASH_STAGE_ENI_LOOKUP_P4_ */
\ No newline at end of file
+#endif /* _DASH_STAGE_ENI_LOOKUP_P4_ */