From 72946403263a0ac056a65f85fd19736d200c8aa4 Mon Sep 17 00:00:00 2001
From: wo_cow <niuqianqian@huawei.com>
Date: Tue, 22 Nov 2022 15:40:11 +0800
Subject: [PATCH] fix bugs for ksliprobe pgsliprobe stackprobe

---
 src/common/common.h                           |  1 +
 src/common/util.c                             | 19 +++++
 .../ebpf.probe/src/ksliprobe/ksliprobe.bpf.c  |  4 +
 .../ebpf.probe/src/pgsliprobe/pgsliprobe.c    | 74 +++++++++----------
 .../src/pgsliprobe/pgsliprobe_bpf.h           |  7 +-
 .../src/stackprobe/conf/stackprobe.conf       |  2 +-
 .../src/stackprobe/stack_bpf/memleak.bpf.c    | 11 +--
 .../extends/ebpf.probe/src/stackprobe/svg.c   | 31 +++++++-
 .../pg_stat.probe/pg_stat_probe.py            |  8 +-
 src/probes/system_infos.probe/system_procs.c  |  2 +-
 10 files changed, 105 insertions(+), 54 deletions(-)

diff --git a/src/common/common.h b/src/common/common.h
index 4ce6a81..49496c4 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -168,6 +168,7 @@ unsigned short ntohs(unsigned short netshort);
 /* get uprobe func offset */
 int get_func_offset(char *proc_name, char *func_name, char *bin_file_path);
 
+char *get_cur_date(void);
 char *get_cur_time(void);
 
 void ip_str(unsigned int family, unsigned char *ip, unsigned char *ip_str, unsigned int ip_str_size);
diff --git a/src/common/util.c b/src/common/util.c
index 65e5d4d..dbe5f75 100644
--- a/src/common/util.c
+++ b/src/common/util.c
@@ -21,6 +21,25 @@
 #include <stdarg.h>
 #include "common.h"
 
+char *get_cur_date(void)
+{
+    /* return date str, ex: 2021/5/17 */
+    static char tm[TM_STR_LEN] = {0};
+    struct tm *tmp_ptr = NULL;
+    time_t t;
+
+    (void)time(&t);
+
+    tmp_ptr = localtime(&t);
+    (void)snprintf(tm,
+        TM_STR_LEN,
+        "%d-%d-%d",
+        (1900 + tmp_ptr->tm_year),
+        (1 + tmp_ptr->tm_mon),
+        tmp_ptr->tm_mday);
+    return tm;
+}
+
 char *get_cur_time(void)
 {
     /* return time str, ex: 2021/5/17 19:56:03 */
diff --git a/src/probes/extends/ebpf.probe/src/ksliprobe/ksliprobe.bpf.c b/src/probes/extends/ebpf.probe/src/ksliprobe/ksliprobe.bpf.c
index bf653f0..6fa64ab 100644
--- a/src/probes/extends/ebpf.probe/src/ksliprobe/ksliprobe.bpf.c
+++ b/src/probes/extends/ebpf.probe/src/ksliprobe/ksliprobe.bpf.c
@@ -374,6 +374,10 @@ static __always_inline void process_rd_msg(u32 tgid, int fd, const char *buf, co
 
 #ifndef KERNEL_SUPPORT_TSTAMP
     csd->start_ts_nsec = ts_nsec;
+#else
+    if (csd->start_ts_nsec == 0) {
+        csd->start_ts_nsec = ts_nsec;
+    }
 #endif
     csd->status = SAMP_READ_READY;
 
diff --git a/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe.c b/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe.c
index 9171985..6b75f06 100644
--- a/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe.c
+++ b/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe.c
@@ -180,7 +180,6 @@ static void *msg_event_receiver(void *arg)
     }
 
     poll_pb(pb, params.period * 1000);
-
     stop = 1;
     return NULL;
 }
@@ -380,6 +379,7 @@ int main(int argc, char **argv)
 {
     int err, ret;
     FILE *fp = NULL;
+    int init = 0;
     struct bpf_link_hash_t *item, *tmp;
 
     err = args_parse(argc, argv, &params);
@@ -409,15 +409,6 @@ int main(int argc, char **argv)
         goto init_err;
     }
 
-    load_args(GET_MAP_FD(pgsli_kprobe, args_map), &params);
-    err = init_conn_mgt_process(GET_MAP_FD(pgsli_kprobe, output));
-    if (err != 0) {
-        fprintf(stderr, "Init connection management process failed.\n");
-        goto init_err;
-    }
-
-    printf("pgsliprobe probe successfully started!\n");
-
     while (!stop) {
         sleep(params.period);
         if (noDependLibssl) {
@@ -426,37 +417,46 @@ int main(int argc, char **argv)
 
         set_bpf_link_inactive();
         if (add_bpf_link_by_search_pids() != SLI_OK) {
-            if (noDependLibssl) {
-                continue;
+            if (!noDependLibssl) {
+                goto init_err;
             }
-            goto init_err;
-        }
-
-        // attach to libssl
-        H_ITER(head, item, tmp) {
-            if (item->v.pid_state == PID_ELF_TOBE_ATTACHED) {
-                UBPF_ATTACH_ONELINK(pgsli_uprobe, SSL_read, item->v.elf_path, SSL_read,
-                    item->v.bpf_link_read, ret);
-                if (ret <= 0) {
-                    fprintf(stderr, "Can't attach function SSL_read at elf_path %s.\n", item->v.elf_path);
-                    goto init_err;
-                }
-                UBPF_RET_ATTACH_ONELINK(pgsli_uprobe, SSL_read, item->v.elf_path, SSL_read,
-                    item->v.bpf_link_read_ret, ret);
-                if (ret <= 0) {
-                    fprintf(stderr, "Can't attach ret function SSL_read at elf_path %s.\n", item->v.elf_path);
-                    goto init_err;
+        } else {
+            // attach to libssl
+            H_ITER(head, item, tmp) {
+                if (item->v.pid_state == PID_ELF_TOBE_ATTACHED) {
+                    UBPF_ATTACH_ONELINK(pgsli_uprobe, SSL_read, item->v.elf_path, SSL_read,
+                        item->v.bpf_link_read, ret);
+                    if (ret <= 0) {
+                        fprintf(stderr, "Can't attach function SSL_read at elf_path %s.\n", item->v.elf_path);
+                        goto init_err;
+                    }
+                    UBPF_RET_ATTACH_ONELINK(pgsli_uprobe, SSL_read, item->v.elf_path, SSL_read,
+                        item->v.bpf_link_read_ret, ret);
+                    if (ret <= 0) {
+                        fprintf(stderr, "Can't attach ret function SSL_read at elf_path %s.\n", item->v.elf_path);
+                        goto init_err;
+                    }
+                    UBPF_ATTACH_ONELINK(pgsli_uprobe, SSL_write, item->v.elf_path, SSL_write,
+                        item->v.bpf_link_write, ret);
+                    if (ret <= 0) {
+                        fprintf(stderr, "Can't attach function SSL_write at elf_path %s.\n", item->v.elf_path);
+                        goto init_err;
+                    }
+                    item->v.pid_state = PID_ELF_ATTACHED;
                 }
-                UBPF_ATTACH_ONELINK(pgsli_uprobe, SSL_write, item->v.elf_path, SSL_write,
-                    item->v.bpf_link_write, ret);
-                if (ret <= 0) {
-                    fprintf(stderr, "Can't attach function SSL_write at elf_path %s.\n", item->v.elf_path);
-                    goto init_err;
-                }
-                item->v.pid_state = PID_ELF_ATTACHED;
             }
+            clear_invalid_bpf_link();
+        }
+        if (init == 0) {
+            load_args(GET_MAP_FD(pgsli_kprobe, args_map), &params);
+            err = init_conn_mgt_process(GET_MAP_FD(pgsli_kprobe, output));
+            if (err != 0) {
+                fprintf(stderr, "Init connection management process failed.\n");
+                goto init_err;
+            }
+            printf("pgsliprobe probe successfully started!\n");
+            init = 1;
         }
-        clear_invalid_bpf_link();
     }
 
 init_err:
diff --git a/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe_bpf.h b/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe_bpf.h
index a2e85b2..a3dfc5d 100644
--- a/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe_bpf.h
+++ b/src/probes/extends/ebpf.probe/src/pgsliprobe/pgsliprobe_bpf.h
@@ -156,7 +156,7 @@ static __always_inline void periodic_report(u64 ts_nsec, struct conn_data_t *con
     if (ts_nsec > conn_data->last_report_ts_nsec &&
         ts_nsec - conn_data->last_report_ts_nsec >= period) {
         // rtt larger than period is considered an invalid value
-        if (conn_data->latency.rtt_nsec < period && conn_data->max.rtt_nsec < period) {
+        if (conn_data->latency.rtt_nsec < period * 2 && conn_data->max.rtt_nsec < period * 2) {
             struct msg_event_data_t msg_evt_data = {0};
             msg_evt_data.tgid = conn_key->tgid;
             msg_evt_data.fd = conn_key->fd;
@@ -224,10 +224,13 @@ static __always_inline void process_rdwr_msg(int fd, const char *buf, int count,
         }
         csd->req_cmd = cmd;
 
+#ifndef KERNEL_SUPPORT_TSTAMP
+        csd->start_ts_nsec = ts_nsec;
+#else
         if (csd->start_ts_nsec == 0) {
             csd->start_ts_nsec = ts_nsec;
         }
-
+#endif
         csd->status = SAMP_READ_READY;
     } else {  // MSG_WRITE
         if (csd->status == SAMP_READ_READY) {
diff --git a/src/probes/extends/ebpf.probe/src/stackprobe/conf/stackprobe.conf b/src/probes/extends/ebpf.probe/src/stackprobe/conf/stackprobe.conf
index f0b3c29..a2edd5d 100644
--- a/src/probes/extends/ebpf.probe/src/stackprobe/conf/stackprobe.conf
+++ b/src/probes/extends/ebpf.probe/src/stackprobe/conf/stackprobe.conf
@@ -12,7 +12,7 @@ flame_name =
     oncpu = true;
     offcpu = false;
     io = false;
-    memleak = true;
+    memleak = false;
 };
 
 application =
diff --git a/src/probes/extends/ebpf.probe/src/stackprobe/stack_bpf/memleak.bpf.c b/src/probes/extends/ebpf.probe/src/stackprobe/stack_bpf/memleak.bpf.c
index e855401..fe2f22b 100644
--- a/src/probes/extends/ebpf.probe/src/stackprobe/stack_bpf/memleak.bpf.c
+++ b/src/probes/extends/ebpf.probe/src/stackprobe/stack_bpf/memleak.bpf.c
@@ -109,20 +109,12 @@ struct bpf_map_def SEC("maps") mmap_allocs = {
 };
 
 struct bpf_map_def SEC("maps") brk_allocs = {
-    .type = BPF_MAP_TYPE_HASH,
+    .type = BPF_MAP_TYPE_LRU_HASH,
     .key_size = sizeof(u32), // tgid
     .value_size = sizeof(struct brk_info_t),
     .max_entries = 1000000,
 };
 
-// allocated memory for the process
-struct bpf_map_def SEC("maps") combined_allocs = {
-    .type = BPF_MAP_TYPE_HASH,
-    .key_size = sizeof(struct stack_id_s),
-    .value_size = sizeof(struct combined_alloc_info_t),
-    .max_entries = 1000,
-};
-
 static __always_inline u64 get_real_start_time()
 {
     struct task_struct* task = (struct task_struct*)bpf_get_current_task();
@@ -185,7 +177,6 @@ static inline void update_statistics(void *ctx, char stackmap_cur, s64 count, st
     }
 }
 
-
 static inline int alloc_exit(void *ctx, u64 addr) {
     u32 tgid = bpf_get_current_pid_tgid() >> INT_LEN;
     struct pid_addr_t pa = {0};
diff --git a/src/probes/extends/ebpf.probe/src/stackprobe/svg.c b/src/probes/extends/ebpf.probe/src/stackprobe/svg.c
index db43aab..78a1cb0 100644
--- a/src/probes/extends/ebpf.probe/src/stackprobe/svg.c
+++ b/src/probes/extends/ebpf.probe/src/stackprobe/svg.c
@@ -156,10 +156,33 @@ static int __create_svg_files(struct stack_svg_s* svg_files, u32 period)
     return 0;
 }
 
+int __mkdir_with_svg_date(const char *svg_dir, char *svg_date_dir, size_t size)
+{
+    size_t len = strlen(svg_dir);
+    const char *day = get_cur_date();
+    if (len <= 1 || len + strlen(day) + 1 >= size) {
+        return -1;
+    }
+
+    (void)snprintf(svg_date_dir, size, "%s/%s", svg_dir, day);
+    if (access(svg_date_dir, F_OK) != 0) {
+        FILE *fp;
+        char command[COMMAND_LEN] = {0};
+        (void)snprintf(command, COMMAND_LEN, "/usr/bin/mkdir -p %s", svg_date_dir);
+        fp = popen(command, "r");
+        if (fp != NULL) {
+            (void)pclose(fp);
+        }
+    }
+
+    return 0;
+}
+
 static int stack_get_next_svg_file(struct stack_svgs_s* svgs, char svg_file[], size_t size, int en_type)
 {
     int next;
     char svg_name[PATH_LEN];
+    char svg_date_dir[PATH_LEN] = {0};
 
     if (svgs->svg_files.files == NULL) {
         return -1;
@@ -176,11 +199,15 @@ static int stack_get_next_svg_file(struct stack_svgs_s* svgs, char svg_file[], s
         svgs->svg_files.files[next] = NULL;
     }
 
+    if (__mkdir_with_svg_date(svgs->svg_dir, svg_date_dir, PATH_LEN) < 0) {
+        return -1;
+    }
+
     svg_name[0] = 0;
-    (void)snprintf(svg_name, PATH_LEN, "%s_%s.svg", svg_params[en_type].file_name, get_cur_time());
+    (void)snprintf(svg_name, PATH_LEN, "%s.svg", get_cur_time());
 
     svg_file[0] = 0;
-    (void)snprintf(svg_file, size, "%s/%s", svgs->svg_dir, svg_name);
+    (void)snprintf(svg_file, size, "%s/%s", svg_date_dir, svg_name);
     __rm_svg(svg_file);
 
     svgs->svg_files.files[next] = strdup(svg_file);
diff --git a/src/probes/extends/python.probe/pg_stat.probe/pg_stat_probe.py b/src/probes/extends/python.probe/pg_stat.probe/pg_stat_probe.py
index 45f80b5..fe8f955 100644
--- a/src/probes/extends/python.probe/pg_stat.probe/pg_stat_probe.py
+++ b/src/probes/extends/python.probe/pg_stat.probe/pg_stat_probe.py
@@ -121,5 +121,11 @@ if __name__ == "__main__":
 
     while True:
         time.sleep(g_period)
-        get_metrics()
+        try:
+            get_metrics()
+        except Exception as e:
+            print("[pg_stat_probe]get metrics failed. Err:" + str(e))
+            stop_conns()
+
+
 
diff --git a/src/probes/system_infos.probe/system_procs.c b/src/probes/system_infos.probe/system_procs.c
index 6812bbc..0e1f8a4 100644
--- a/src/probes/system_infos.probe/system_procs.c
+++ b/src/probes/system_infos.probe/system_procs.c
@@ -31,7 +31,7 @@
 #define PROC_START_TIME_CMD "/usr/bin/cat /proc/%s/stat | awk '{print $22}'"
 #define PROC_CMDLINE_CMD    "/proc/%s/cmdline"
 #define PROC_FD             "/proc/%s/fd"
-#define PROC_FD_CNT_CMD     "/usr/bin/ls -l /proc/%s/fd | wc -l"
+#define PROC_FD_CNT_CMD     "/usr/bin/ls -l /proc/%s/fd 2>/dev/null | wc -l 2>/dev/null"
 #define PROC_IO             "/proc/%s/io"
 #define PROC_IO_CMD         "/usr/bin/cat /proc/%s/io"
 #define PROC_SMAPS          "/proc/%s/smaps_rollup"
-- 
2.28.0.windows.1

