Fixed #TASK_QT-13496 数据结构暂时改造完成。

ebpftracer/ebpf/netflow/kflowd.bpf.c

@@ -46,20 +46,6 @@ struct {
     __uint(value_size, sizeof(int));
 } perfevent_records SEC(".maps");
 
-struct {
-    __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAP_RECORDS_MAX);
-    __type(key, __u64);
-    __type(value, struct RECORD_FS);
-} hash_records SEC(".maps");
-
-struct {
-    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
-    __uint(max_entries, 1);
-    __type(key, int);
-    __type(value, struct RECORD_FS);
-} heap_record_fs SEC(".maps");
-
 struct {
     __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
     __uint(max_entries, 1);
@@ -67,10 +53,17 @@ struct {
     __type(value, struct RECORD_SOCK);
 } heap_record_sock SEC(".maps");
 
+// struct {
+//     __uint(type, BPF_MAP_TYPE_LRU_HASH);
+//     __uint(max_entries, MAP_SOCKS_MAX);
+//     __type(key, __u64);
+//     __type(value, struct SOCK_INFO);
+// } hash_socks SEC(".maps");
+
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __uint(max_entries, MAP_SOCKS_MAX);
-    __type(key, __u64);
+    __type(key, struct SOCK_TUPLE);
     __type(value, struct SOCK_INFO);
 } hash_socks SEC(".maps");
 
@@ -81,12 +74,19 @@ struct {
     __type(value, struct SOCK_INFO);
 } heap_sock SEC(".maps");
 
+// struct {
+//     __uint(type, BPF_MAP_TYPE_LRU_HASH);
+//     __uint(max_entries, MAP_SOCKS_MAX);
+//     __type(key, struct SOCK_TUPLE);
+//     __type(value, __u64);
+// } hash_tuples SEC(".maps");
+
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __uint(max_entries, MAP_SOCKS_MAX);
-    __type(key, struct SOCK_TUPLE);
-    __type(value, __u64);
-} hash_tuples SEC(".maps");
+    __type(key, __u64);
+    __type(value, struct SOCK_TUPLE);
+} hash_key_tuples SEC(".maps");
 
 struct {
     __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
@@ -95,51 +95,6 @@ struct {
     __type(value, struct SOCK_TUPLE);
 } heap_tuple SEC(".maps");
 
-// struct {
-//     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-//     __uint(max_entries, MAP_XFILES_MAX);
-//     __type(key, char[CW_NET_TASK_COMM_LEN]);
-//     __type(value, struct XFILES);
-// } hash_xfiles SEC(".maps");
-
-// struct {
-//     __uint(type, BPF_MAP_TYPE_QUEUE);
-//     __uint(max_entries, MAP_SOCKS_MAX);
-//     __type(value, __u64[2]);
-// } queue_socks SEC(".maps");
-
-struct {
-    __uint(type, BPF_MAP_TYPE_ARRAY);
-    __uint(max_entries, 1);
-    __type(key, int);
-    __type(value, struct STATS);
-} stats SEC(".maps");
-
-// #define bpf_core_read(dst, sz, src)	bpf_probe_read_kernel(dst, sz, (const void *)__builtin_preserve_access_index(src))
-
-// #define ___concat(a, b) a ## b
-// #define ___apply(fn, n) ___concat(fn, n)
-// #define ___nth(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, __11, N, ...) N
-// #define ___narg(...) ___nth(_, ##__VA_ARGS__, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0)
-// #define ___arrow(...) ___apply(___arrow, ___narg(__VA_ARGS__))(__VA_ARGS__)
-// #define ___type(...) typeof(___arrow(__VA_ARGS__))
-
-// #define ___empty(...) ___nth(_, ##__VA_ARGS__, N, N, N, N, N, N, N, N, N, N, 0)
-// #define ___core_read(fn, fn_ptr, dst, src, a, ...)			    \
-// 	___apply(___core_read, ___empty(__VA_ARGS__))(fn, fn_ptr, dst,	    \
-// 						      src, a, ##__VA_ARGS__)
-
-// #define BPF_CORE_READ_INTO(dst, src, a, ...) ({				    \
-// 	___core_read(bpf_core_read, bpf_core_read,			    \
-// 		     dst, (src), a, ##__VA_ARGS__)			    \
-// })
-
-// #define BPF_CORE_READ(src, a, ...) ({					    \
-// 	___type((src), a, ##__VA_ARGS__) __r;				    \
-// 	BPF_CORE_READ_INTO(&__r, (src), a, ##__VA_ARGS__);		    \
-// 	__r;								    \
-// })
-
 /* glabal variables shared with userspace */
 // const volatile __u64 ts_start;
 // const volatile __u32 agg_events_max;
@@ -322,7 +277,7 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
     __u8                tcp_state;
     char               *func;
     __u64               key;
-    __u64               key_alt;
+    __u64               key_alt = 0;
     __u32               zero = 0;
     __u32               cnt;
     
@@ -378,9 +333,11 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
         stuple->lport = args_sport;
         stuple->rport = args_dport;
         stuple->proto = IPPROTO_TCP;
-        if (bpf_map_update_elem(&hash_tuples, stuple, &key, BPF_ANY))
+        
+        // if (bpf_map_update_elem(&hash_tuples, stuple, &key, BPF_ANY))
+        //     bpf_printk("WARNING: Failed to update client/server stuple for key %lx and pid %u\n", key, pid);
+        if (bpf_map_update_elem(&hash_key_tuples,&key,stuple, BPF_ANY))
             bpf_printk("WARNING: Failed to update client/server stuple for key %lx and pid %u\n", key, pid);
-
         /* get old and new tcp state */
         int args_oldstate;
         int args_newstate;
@@ -391,8 +348,8 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
         if (tcp_state_old == TCP_SYN_RECV && tcp_state == TCP_ESTABLISHED) {
             /* check if alternate key from application message exists already */
             // key_alt = crc64(0, (const u8 *)stuple, sizeof(*stuple));
-            key_alt = 0;
-            sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
+            // sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
+            sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
             if (!sinfo) {
                 sinfo = bpf_map_lookup_elem(&heap_sock, &zero);
                 if (!sinfo) {
@@ -443,7 +400,8 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
             sinfo->tx_event[0] = 1;
             sinfo->tx_flags_map[0] = TCP_SYN | TCP_ACK;
             sinfo->tx_flags_map_cnt = 1;
-            if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+            // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+            if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
                 // if (debug_proc(NULL, NULL))
                 // bpf_printk("Prepared %s server socket for pid %u\n",
                 //                sinfo->app_msg.cnt ? "new tcp" : "tcp application", pid);
@@ -505,26 +463,28 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
             /* calculate alternate key for tuple since no kernel socket hash at this point */
             sinfo->app_msg.cnt = 0;
             // key_alt = crc64(0, (const u8 *)stuple, sizeof(*stuple));
-            key_alt = 0;
-            if (!bpf_map_update_elem(&hash_socks, &key_alt, sinfo, BPF_ANY)) {
+            // if (!bpf_map_update_elem(&hash_socks, &key_alt, sinfo, BPF_ANY)) {
+            if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
                 // if (debug_proc(sinfo->comm, NULL))
-                    bpf_printk("Prepared new tcp client socket for alt key %lx and pid %u\n", key_alt, pid);
+                    // bpf_printk("Prepared new tcp client socket for alt key %lx and pid %u\n", key_alt, pid);
+                    ;
             } else
                 bpf_printk("WARNING: Failed to prepare new tcp client socket for alt key %lx and pid %u\n", key_alt,
                            pid);
         } else if (tcp_state_old == TCP_SYN_SENT && tcp_state == TCP_ESTABLISHED) {
             /* get alternate key */
             // key_alt = crc64(0, (const u8 *)stuple, sizeof(*stuple));
-            key_alt = 0;
-            sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
+            // sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
+            sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
             if (!sinfo || (struct cw_net_sock*)sinfo->sock != sock) {
                 /* try again without lport */
                 u16 lport = stuple->lport;
                 stuple->lport = 0;
                 // key_alt = crc64(0, (const u8 *)stuple, sizeof(*stuple));
-                key_alt = 0;
+                // stuple->lport = lport;
+                // sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
+                sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
                 stuple->lport = lport;
-                sinfo = bpf_map_lookup_elem(&hash_socks, &key_alt);
                 if (!sinfo || (struct cw_net_sock*)sinfo->sock != sock) {
                     bpf_printk("WARNING: Failed lookup to add tcp client socket for alt key %lx and pid %u\n", key_alt,
                                pid);
@@ -533,7 +493,8 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
             }
             sinfo->state = tcp_state;
             /* add new tcp client socket */
-            if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+            // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+            if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
                 // if (debug_proc(sinfo->comm, NULL))
                     bpf_printk("Added new tcp client socket for alt key %lx, key %lx and pid %u\n", key_alt, key,
                                sinfo->pid);
@@ -542,7 +503,8 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
                            key_alt, key, sinfo->pid);
         } else if ((tcp_state_old == TCP_LAST_ACK && tcp_state == TCP_CLOSE) ||
                    (tcp_state_old == TCP_FIN_WAIT2 && tcp_state == TCP_CLOSE)) {
-            sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+            // sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+            sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
             if (!sinfo || (struct cw_net_sock*)sinfo->sock != sock) {
                 bpf_printk("WARNING: Failed lookup to delete tcp socket for key %lx, lport %u and pid %u", key,
                            stuple->lport, pid);
@@ -551,7 +513,8 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
             /* submit final record and delete closed client and server sockets */
             sinfo->state = tcp_state;
             submit_sock_record(ctx, sinfo);
-            if (bpf_map_delete_elem(&hash_socks, &key))
+            // if (bpf_map_delete_elem(&hash_socks, &key))
+            if (bpf_map_delete_elem(&hash_socks, stuple))
                 // bpf_printk("WARNING: Failed to delete %s socket for key %lx and pid %u\n", GET_ROLE_STR(sinfo->role),
                         //    key, sinfo->pid);
                 bpf_printk("WARNING: Failed to delete\n");
@@ -577,7 +540,14 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
             // 错误处理  
         }
         key = KEY_SOCK(skc_hash);
-        sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+
+        stuple = bpf_map_lookup_elem(&hash_key_tuples, &key);
+        if (!stuple) {
+            bpf_printk("WARNING: Failed to allocate new tuple for pid %u\n", pid);
+            return 0;
+        }
+        // sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+        sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
         if (!sinfo || (struct cw_net_sock*)sinfo->sock != sock)
             return 0;
         sinfo->sock = (uint64_t)sock;
@@ -610,21 +580,24 @@ static __always_inline int handle_tcp_event(void *ctx, const struct SOCK_EVENT_I
         // }
 
         /* set tuple */
-        stuple = bpf_map_lookup_elem(&heap_tuple, &zero);
-        if (!stuple) {
-            bpf_printk("WARNING: Failed to allocate new tuple for pid %u\n", pid);
-            return 0;
-        }
-        bpf_probe_read_kernel(stuple->laddr, sizeof(stuple->laddr), sinfo->laddr);
-        bpf_probe_read_kernel(stuple->raddr, sizeof(stuple->raddr), sinfo->raddr);
-        stuple->lport = sinfo->lport;
-        stuple->rport = sinfo->rport;
-        stuple->proto = IPPROTO_TCP;
-        if (bpf_map_update_elem(&hash_tuples, stuple, &key, BPF_ANY))
+        // stuple = bpf_map_lookup_elem(&heap_tuple, &zero);
+        // if (!stuple) {
+        //     bpf_printk("WARNING: Failed to allocate new tuple for pid %u\n", pid);
+        //     return 0;
+        // }
+        // bpf_probe_read_kernel(stuple->laddr, sizeof(stuple->laddr), sinfo->laddr);
+        // bpf_probe_read_kernel(stuple->raddr, sizeof(stuple->raddr), sinfo->raddr);
+        // stuple->lport = sinfo->lport;
+        // stuple->rport = sinfo->rport;
+        // stuple->proto = IPPROTO_TCP;
+        // if (bpf_map_update_elem(&hash_tuples, stuple, &key, BPF_ANY))
+        //     bpf_printk("WARNING: Failed to update tcp server stuple for key %lx and pid %u\n", key, pid);
+        if (bpf_map_update_elem(&hash_key_tuples,&key,stuple, BPF_ANY))
             bpf_printk("WARNING: Failed to update tcp server stuple for key %lx and pid %u\n", key, pid);
 
         /* update hash tables */
-        if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
             // if (debug_proc(sinfo->comm, NULL))
             //     bpf_printk("Added new tcp server socket for key %lx, rport %u and pid %u\n", key, sinfo->rport, pid);
             ;
@@ -689,6 +662,7 @@ static __always_inline int handle_tcp_packet(void* ctx, struct cw_net_sock *sock
     // KPROBE_SWITCH(MONITOR_SOCK);
     struct cw_net_task_struct *task = (struct cw_net_task_struct *)bpf_get_current_task();
     struct SOCK_INFO   *sinfo;
+    struct SOCK_TUPLE  *stuple;
     struct SOCK_QUEUE   sq = {0};
     struct STATS       *s = NULL;
     __u8                tcp_flags = 0;
@@ -718,7 +692,13 @@ static __always_inline int handle_tcp_packet(void* ctx, struct cw_net_sock *sock
         // 错误处理  
     }
     key = KEY_SOCK(skc_hash);
-    sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+    stuple = bpf_map_lookup_elem(&hash_key_tuples, &key);
+    if (!stuple) {
+        bpf_printk("WARNING: stuple is invalid\n");
+        return 0;
+    } 
+    // sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+    sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
     const char* skb_head;
     unsigned int skb_end;
     unsigned int skb_len;
@@ -981,7 +961,7 @@ static __always_inline int handle_tcp_packet(void* ctx, struct cw_net_sock *sock
             bpf_probe_read_kernel(&srtt_us_tmp, sizeof(srtt_us_tmp), &(tcp_sock->srtt_us));
             sinfo->rtt = srtt_us_tmp * 1000 / 8;
         }
-        if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
             // if (debug_proc(sinfo->comm, NULL))
             //     bpf_printk("Updated tcp %s flags of socket %lx for pid %u", isrx ? "rx" : "tx", key, sinfo->pid);
             sq.key = key;
@@ -1089,6 +1069,7 @@ static __always_inline int handle_udp_event(void *ctx, const struct SOCK_EVENT_I
     struct udphdr          *udphdr;
     struct SOCK_INFO       *sinfo;
     struct SOCK_QUEUE       sq = {0};
+    struct SOCK_TUPLE  *stuple;
     struct STATS           *s;
     __u16                   gso_segs;
     char                    comm[CW_NET_TASK_COMM_LEN] = {0};
@@ -1228,8 +1209,13 @@ static __always_inline int handle_udp_event(void *ctx, const struct SOCK_EVENT_I
         // 错误处理  
     }
     key = KEY_SOCK(skc_hash);
-    sinfo = bpf_map_lookup_elem(&hash_socks, &key);
-    s = bpf_map_lookup_elem(&stats, &zero);
+    // sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+    stuple = bpf_map_lookup_elem(&hash_key_tuples, &key);
+    if (!stuple) {
+        bpf_printk("WARNING: Failed to allocate new tuple for pid %u\n", pid);
+        return 0;
+    }
+    sinfo = bpf_map_lookup_elem(&hash_socks, stuple);
     ts_now = bpf_ktime_get_ns();
     if (sinfo && (struct cw_net_sock*)sinfo->sock == sock) {
         /* update existing udp socket */
@@ -1284,7 +1270,8 @@ static __always_inline int handle_udp_event(void *ctx, const struct SOCK_EVENT_I
                 /* export record on max application messages */
                 if (sinfo->app_msg.cnt >= APP_MSG_MAX) {
                     submit_sock_record(ctx, sinfo);
-                    if (bpf_map_delete_elem(&hash_socks, &key))
+                    // if (bpf_map_delete_elem(&hash_socks, &key))
+                    if (bpf_map_delete_elem(&hash_socks, stuple))
                         // bpf_printk("WARNING: Failed to delete %s socket for key %lx and pid %u\n",
                         //            GET_ROLE_STR(sinfo->role), key, sinfo->pid);
                         bpf_printk("bpf_map_delete_elem failed\n");
@@ -1298,7 +1285,8 @@ static __always_inline int handle_udp_event(void *ctx, const struct SOCK_EVENT_I
             } else
                 bpf_printk("WARNING: Failed to capture dns application message #%u\n", sinfo->app_msg.cnt);
         }
-        if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
             // if (debug_proc(sinfo->comm, NULL))
                 // bpf_printk("Updated %s socket %lx for pid %u", GET_ROLE_STR(sinfo->role), key, pid);
             // bpf_printk("bpf_map_update_elem success \n");
@@ -1452,7 +1440,8 @@ static __always_inline int handle_udp_event(void *ctx, const struct SOCK_EVENT_I
                 sinfo->role = ROLE_UDP_SERVER;
         }
 
-        if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+        if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
             // if (debug_proc(sinfo->comm, NULL))
                 // bpf_printk("Added new %s socket %lx for pid %u", GET_ROLE_STR(sinfo->role), key, pid);
             bpf_printk("bpf_map_update_elem success\n");
@@ -1725,16 +1714,16 @@ int handle_skb(struct __sk_buff *skb) {
     stuple->lport = lport;
     stuple->rport = rport;
     stuple->proto = proto;
-    pkey = bpf_map_lookup_elem(&hash_tuples, stuple);
-    if (pkey) {
-        // bpf_probe_read_kernel(&key, sizeof(key), pkey);
-        key = *pkey;
-        sinfo = bpf_map_lookup_elem(&hash_socks, &key);
-        if (!sinfo) {
-            bpf_printk("WARNING: Failed to lookup tcp socket key %lx for lport %u and rport %u\n", key, lport, rport);
-            return skb->len;
-        }
-    }
+    // pkey = bpf_map_lookup_elem(&hash_tuples, stuple);
+    // if (pkey) {
+    //     // bpf_probe_read_kernel(&key, sizeof(key), pkey);
+    //     key = *pkey;
+    //     sinfo = bpf_map_lookup_elem(&hash_socks, &key);
+    //     if (!sinfo) {
+    //         bpf_printk("WARNING: Failed to lookup tcp socket key %lx for lport %u and rport %u\n", key, lport, rport);
+    //         return skb->len;
+    //     }
+    // }
     // return 0;
     if (!sinfo) {
         if(!isrx)
@@ -1803,7 +1792,8 @@ int handle_skb(struct __sk_buff *skb) {
     // }
     // else
     //     return skb->len;
-    if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+    // if (!bpf_map_update_elem(&hash_socks, &key, sinfo, BPF_ANY)) {
+    if (!bpf_map_update_elem(&hash_socks, stuple, sinfo, BPF_ANY)) {
         // if (debug_proc(sinfo->comm, NULL))
         //     bpf_printk("Captured payload for %s socket %lx and pid %u", GET_ROLE_STR(sinfo->role), key, sinfo->pid);
         bpf_printk("success\n");