From 28fe5004ce76c55ab254a3b5c4f65a2dfce8f8f8 Mon Sep 17 00:00:00 2001
From: wuchangye <wuchangye@huawei.com>
Date: Tue, 23 Jul 2024 12:20:13 +0800
Subject: [PATCH] update metric_key with direction & dst_port

Signed-off-by: wuchangye <wuchangye@huawei.com>
---
 bpf/kmesh/probes/metrics.h | 40 ++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/bpf/kmesh/probes/metrics.h b/bpf/kmesh/probes/metrics.h
index 3ad57ed54..fd8a7114d 100644
--- a/bpf/kmesh/probes/metrics.h
+++ b/bpf/kmesh/probes/metrics.h
@@ -9,6 +9,8 @@
 struct metric_key {
     struct ip_addr src_ip;
     struct ip_addr dst_ip;
+    __u32 direction;
+    __u32 dst_port;
 };
 
 struct metric_data {
@@ -34,20 +36,34 @@ struct {
     __uint(max_entries, RINGBUF_SIZE);
 } map_of_metric_notify SEC(".maps");
 
-static inline void construct_metric_key(struct bpf_sock *sk, struct metric_key *key)
+static inline void construct_metric_key(struct bpf_sock *sk, __u8 direction, struct metric_key *key)
 {
     bpf_memset(key, 0, sizeof(struct metric_key));
-    if (sk->family == AF_INET) {
-        key->src_ip.ip4 = sk->src_ip4;
-        key->dst_ip.ip4 = sk->dst_ip4;
+
+    key->direction = direction;
+    if (direction == OUTBOUND) {
+        if (sk->family == AF_INET) {
+            key->src_ip.ip4 = sk->src_ip4;
+            key->dst_ip.ip4 = sk->dst_ip4;
+        } else {
+            bpf_memcpy(key->src_ip.ip6, sk->src_ip6, IPV6_ADDR_LEN);
+            bpf_memcpy(key->dst_ip.ip6, sk->dst_ip6, IPV6_ADDR_LEN);
+        }
+        key->dst_port = bpf_ntohl(sk->dst_port);
     } else {
-        bpf_memcpy(key->src_ip.ip6, sk->src_ip6, IPV6_ADDR_LEN);
-        bpf_memcpy(key->dst_ip.ip6, sk->dst_ip6, IPV6_ADDR_LEN);
+        if (sk->family == AF_INET) {
+            key->src_ip.ip4 = sk->dst_ip4;
+            key->dst_ip.ip4 = sk->src_ip4;
+        } else {
+            bpf_memcpy(key->src_ip.ip6, sk->dst_ip6, IPV6_ADDR_LEN);
+            bpf_memcpy(key->dst_ip.ip6, sk->src_ip6, IPV6_ADDR_LEN);
+        }
+        key->dst_port = sk->src_port;
     }
     return;
 }
 
-static inline void report_metrics(struct bpf_sock *sk)
+static inline void report_metrics(struct metric_key *mk)
 {
     struct metric_key *key = bpf_ringbuf_reserve(&map_of_metric_notify, sizeof(struct metric_key), 0);
     if (!key) {
@@ -55,7 +71,7 @@ static inline void report_metrics(struct bpf_sock *sk)
         return;
     }
 
-    construct_metric_key(sk, key);
+    bpf_memcpy(key, mk, sizeof(struct metric_key));
     bpf_ringbuf_submit(key, 0);
     return;
 }
@@ -67,7 +83,7 @@ metric_on_connect(struct bpf_sock *sk, struct bpf_tcp_sock *tcp_sock, struct soc
     struct metric_data data = {0};
     struct metric_data *metric = NULL;
 
-    construct_metric_key(sk, &key);
+    construct_metric_key(sk, storage->direction, &key);
     metric = (struct metric_data *)bpf_map_lookup_elem(&map_of_metrics, &key);
     if (!metric) {
         data.conn_open++;
@@ -83,7 +99,7 @@ metric_on_connect(struct bpf_sock *sk, struct bpf_tcp_sock *tcp_sock, struct soc
     metric->conn_open++;
     metric->direction = storage->direction;
 notify:
-    report_metrics(sk);
+    report_metrics(&key);
     return;
 }
 
@@ -94,7 +110,7 @@ metric_on_close(struct bpf_sock *sk, struct bpf_tcp_sock *tcp_sock, struct sock_
     struct metric_data data = {0};
     struct metric_data *metric = NULL;
 
-    construct_metric_key(sk, &key);
+    construct_metric_key(sk, storage->direction, &key);
     metric = (struct metric_data *)bpf_map_lookup_elem(&map_of_metrics, &key);
     if (!metric) {
         // connect failed
@@ -113,7 +129,7 @@ metric_on_close(struct bpf_sock *sk, struct bpf_tcp_sock *tcp_sock, struct sock_
     metric->sent_bytes += tcp_sock->delivered;
     metric->received_bytes += tcp_sock->bytes_received;
 notify:
-    report_metrics(sk);
+    report_metrics(&key);
     return;
 }