From cb0c86801bc4c221ae79a284619879dc98a5464d Mon Sep 17 00:00:00 2001
From: hantwofish <hankangkang5@huawei.com>
Date: Sat, 29 Jun 2024 14:46:32 +0800
Subject: [PATCH] virtio: update g_rule_port by reg_ring_type enum

---
 src/lstack/core/lstack_port_map.c    | 25 ++++++++++++-------------
 src/lstack/include/lstack_port_map.h |  4 ++--
 src/lstack/netif/lstack_vdev.c       | 15 +++++++++++++++
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/src/lstack/core/lstack_port_map.c b/src/lstack/core/lstack_port_map.c
index e5008b3..5439394 100644
--- a/src/lstack/core/lstack_port_map.c
+++ b/src/lstack/core/lstack_port_map.c
@@ -15,29 +15,28 @@
 #include "lstack_port_map.h"
 
 #define PORT_MAP_UNIX_TCP_PORT_MAX 65535
-#define PORT_MAP_EIGHT_BIT 8
 
-static uint8_t g_rule_port[(PORT_MAP_UNIX_TCP_PORT_MAX + 1) / PORT_MAP_EIGHT_BIT]; // 8k byte
+static uint32_t g_rule_port[PORT_MAP_UNIX_TCP_PORT_MAX] = {0};
 static pthread_mutex_t g_rule_map_mutex = PTHREAD_MUTEX_INITIALIZER;
 
-void port_map_set(uint32_t modBit, int setVal)
+void port_map_mod(uint16_t port, uint16_t flag)
 {
     pthread_mutex_lock(&g_rule_map_mutex);
-    g_rule_port[modBit / PORT_MAP_EIGHT_BIT] &= ~(1 << (modBit % PORT_MAP_EIGHT_BIT));
-    g_rule_port[modBit / PORT_MAP_EIGHT_BIT] |= (setVal << (modBit % PORT_MAP_EIGHT_BIT));
+    if (flag == 0) {
+        g_rule_port[port]--;
+    } else {
+        g_rule_port[port]++;
+    }
     pthread_mutex_unlock(&g_rule_map_mutex);
 }
 
-int port_map_get(int bit_index)
+uint16_t port_map_get(uint16_t port)
 {
-    int bit_val = 0;
-    int byte_index = bit_index / PORT_MAP_EIGHT_BIT;
-    int bit_offset = bit_index % PORT_MAP_EIGHT_BIT;
-    uint8_t mask = 1 << bit_offset;
+    uint16_t val = 0;
     pthread_mutex_lock(&g_rule_map_mutex);
-    if ((g_rule_port[byte_index] & mask) != 0) {
-        bit_val = 1;
+    if (g_rule_port[port] > 0) {
+        val = 1;
     }
     pthread_mutex_unlock(&g_rule_map_mutex);
-    return bit_val;
+    return val;
 }
\ No newline at end of file
diff --git a/src/lstack/include/lstack_port_map.h b/src/lstack/include/lstack_port_map.h
index 24ef53a..80955dc 100644
--- a/src/lstack/include/lstack_port_map.h
+++ b/src/lstack/include/lstack_port_map.h
@@ -14,7 +14,7 @@
 
 #include <stdint.h>
 
-void port_map_set(uint32_t modBit, int setVal);
-int port_map_get(int bit_index);
+void port_map_mod(uint16_t port, uint16_t flag);
+uint16_t port_map_get(uint16_t port);
 
 #endif
\ No newline at end of file
diff --git a/src/lstack/netif/lstack_vdev.c b/src/lstack/netif/lstack_vdev.c
index cb3953e..75e4acc 100644
--- a/src/lstack/netif/lstack_vdev.c
+++ b/src/lstack/netif/lstack_vdev.c
@@ -35,6 +35,8 @@
 #include "lstack_lwip.h"
 #include "lstack_flow.h"
 #include "lstack_vdev.h"
+#include "lstack_port_map.h"
+#include "lstack_virtio.h"
 
 /* INUSE_TX_PKTS_WATERMARK < VDEV_RX_QUEUE_SZ;
  * USE_RX_PKTS_WATERMARK < FREE_RX_QUEUE_SZ.
@@ -197,6 +199,19 @@ int32_t vdev_reg_xmit(enum reg_ring_type type, struct gazelle_quintuple *qtuple)
         return -1;
     }
 
+    uint16_t local_port = ntohs(qtuple->src_port);
+    if (get_global_cfg_params()->flow_bifurcation && get_global_cfg_params()->is_primary) {
+        if (type == REG_RING_TCP_LISTEN_CLOSE || type == REG_RING_UDP_BIND_CLOSE) { // add enum type in reg_sock.h
+            port_map_mod(local_port, 0);
+        } else if (type == REG_RING_TCP_LISTEN || type == REG_RING_UDP_BIND) {
+            port_map_mod(local_port, 1);
+        } else if (type == REG_RING_TCP_CONNECT_CLOSE) {
+            port_map_mod(local_port, 0);
+        } else if (type == REG_RING_TCP_CONNECT) {
+            port_map_mod(local_port, 1);
+        }
+    }
+
     if (!use_ltran() && get_global_cfg_params()->tuple_filter) {
         if (type == REG_RING_TCP_LISTEN_CLOSE) {
             if (get_global_cfg_params()->is_primary) {
-- 
2.33.0

