Merge pull request #627 from nlgwcy/ipv6-1

update metric_key with direction & dst_port
kmesh-net · Jul 23, 2024 · 064614b · 064614b
2 parents ecaee92 + 28fe500
commit 064614b
Showing 1 changed file with 28 additions and 12 deletions.
diff --git 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,28 +36,42 @@ 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) {
         BPF_LOG(ERR, PROBE, "report_metrics bpf_ringbuf_reserve failed\n");
         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;
 }