edriver-rust: patch ipv6

chriskaliX · Feb 19, 2024 · d40feae · d40feae
1 parent c1b21b4
commit d40feae
Show file tree

Hide file tree

Showing 8 changed files with 9,928 additions and 9,672 deletions.
diff --git a/plugins/edriver-rust/src/bpf/common/common.h b/plugins/edriver-rust/src/bpf/common/common.h
@@ -23,6 +23,34 @@ static inline struct hds_context init_context(void *ctx, int dt)
     return c; 
 }
 
+static inline volatile unsigned char get_sock_state(struct sock *sock)
+{
+    volatile unsigned char sk_state_own_impl;
+    bpf_probe_read(
+        (void *) &sk_state_own_impl, sizeof(sk_state_own_impl), (const void *) &sock->sk_state);
+    return sk_state_own_impl;
+}
+
+static inline struct ipv6_pinfo *get_inet_pinet6(struct inet_sock *inet)
+{
+    struct ipv6_pinfo *pinet6_own_impl;
+    bpf_probe_read(&pinet6_own_impl, sizeof(pinet6_own_impl), &inet->pinet6);
+    return pinet6_own_impl;
+}
+
+
+static inline struct ipv6_pinfo *inet6_sk_own_impl(struct sock *__sk, struct inet_sock *inet)
+{
+    volatile unsigned char sk_state_own_impl;
+    sk_state_own_impl = get_sock_state(__sk);
+
+    struct ipv6_pinfo *pinet6_own_impl;
+    pinet6_own_impl = get_inet_pinet6(inet);
+
+    bool sk_fullsock = (1 << sk_state_own_impl) & ~(TCPF_TIME_WAIT | TCPF_NEW_SYN_RECV);
+    return sk_fullsock ? pinet6_own_impl : NULL;
+}
+
 /* notice: char * to void * */
 static __noinline int do_u32toa(uint32_t v, void *s, int l)
 {
@@ -200,6 +228,21 @@ static __always_inline int get_sock_v4(struct sock *sk, struct hds_socket_info *
     return 0;
 }
 
+static __always_inline int get_sock_v6(struct sock *sk, struct hds_socket_info_v6 *sinfo)
+{
+    struct inet_sock *inet = (struct inet_sock *)sk;
+    struct ipv6_pinfo *inet6 = inet6_sk_own_impl(sk, inet);
+    struct in6_addr addr = {};
+    addr = BPF_CORE_READ(sk, __sk_common.skc_v6_rcv_saddr);
+    if (ipv6_addr_any(&addr))
+        addr = BPF_CORE_READ(inet6, saddr);
+    sinfo->local_address = BPF_CORE_READ(sk, __sk_common.skc_v6_daddr);
+    sinfo->local_port = BPF_CORE_READ(inet, inet_dport);
+    sinfo->remote_address = addr;
+    sinfo->remote_port = BPF_CORE_READ(inet, inet_sport);
+    return 0;
+}
+
 /* ===== END ===== */
 
 
@@ -247,11 +290,11 @@ static __always_inline void *get_path(struct path *path)
         d_name = BPF_CORE_READ(dentry, d_name);
         off = buf_off - (d_name.len + 1);
         sz = 0;
-        /* off check */
-        if (off > buf_off)
-            break;
         /* size check */
-        sz = bpf_probe_read_str(&(cache->buf[off & MID_PERCPU_MASK]), (d_name.len + 1) & MID_PERCPU_MASK, (void *)d_name.name);
+        off = off & MAX_PERCPU_MASK;
+        if (off > MAX_PERCPU_MASK - MAX_STRING_SIZE)
+            break;
+        sz = bpf_probe_read_str(&(cache->buf[off]), (d_name.len + 1) & MAX_STRING_MASK, (void *)d_name.name);
         if (!sz)
             break;
         cache->buf[(buf_off - 1) & MAX_PERCPU_MASK] = '/';

diff --git a/plugins/edriver-rust/src/bpf/common/consts.h b/plugins/edriver-rust/src/bpf/common/consts.h
@@ -4,6 +4,7 @@
 /* consts: constrains */
 #define TASK_COMM_LEN               (16)
 #define MAX_STRING_SIZE             (256)
+#define MAX_STRING_MASK             (MAX_STRING_SIZE - 1)
 #define MAX_STR_ARR_ELEM            (32)
 #define MAX_PATH_COMPONENTS         (16)
 #define MAX_PATH_COMPONENTS_SIM     (10)

diff --git a/plugins/edriver-rust/src/bpf/common/edriver.h b/plugins/edriver-rust/src/bpf/common/edriver.h
@@ -61,15 +61,22 @@ int rtp__process_exec(struct bpf_raw_tracepoint_args *ctx)
     SBT_CHAR((&c), get_task_tty(task));
     /* pwd */
     struct path pwd = BPF_CORE_READ(task, fs, pwd);
-    SBT_CHAR((&c), get_path(__builtin_preserve_access_index(&pwd)));
+    void *pwd_ptr = get_path(__builtin_preserve_access_index(&pwd));
+    SBT_CHAR((&c), pwd_ptr);
     /* stdin */
     SBT_CHAR((&c), get_fd(task, 0));
     /* stdout */
     SBT_CHAR((&c), get_fd(task, 1));
     /* exe */
     struct path exe = BPF_CORE_READ(task, mm, exe_file, f_path);
-    SBT_CHAR((&c), get_path(__builtin_preserve_access_index(&exe)));
-    SBT((&c), &proc_i->sinfo, sizeof(struct hds_socket_info));
+    void *exe_ptr = get_path(__builtin_preserve_access_index(&exe));
+    SBT_CHAR((&c), exe_ptr);
+    /* socket info */
+    SBT((&c), &proc_i->family, S_U16);
+    if (proc_i->family == AF_INET6)
+        SBT((&c), &proc_i->sinfo_v6, sizeof(struct hds_socket_info_v6));
+    else if (proc_i->family == AF_INET)
+        SBT((&c), &proc_i->sinfo, sizeof(struct hds_socket_info));
     SBT_CHAR((&c), &proc_i->pidtree);
 
     return report_event(&c);
@@ -109,13 +116,17 @@ static struct proc_info *proc_info_init(struct task_struct *task)
     struct sock *sk = proc_socket_info(task, &proc_i->socket_pid);
     if (!sk) {
         proc_i->socket_pid = 0;
-    } else {
-        struct hds_socket_info sinfo = {};
-        sinfo.family = BPF_CORE_READ(sk, sk_family);
-        if (sinfo.family == AF_INET)
+    } else {        
+        proc_i->family = BPF_CORE_READ(sk, sk_family);
+        if (proc_i->family == AF_INET) {
+            struct hds_socket_info sinfo = {};
             get_sock_v4(sk, &sinfo);
-        // else if (sinfo.family == AF_INET6)
-        proc_i->sinfo = sinfo;
+            proc_i->sinfo = sinfo;
+        } else if (proc_i->family == AF_INET6) {
+            struct hds_socket_info_v6 sinfo_v6 = {};
+            get_sock_v6(sk, &sinfo_v6);
+            proc_i->sinfo_v6 = sinfo_v6;
+        }
     }
     /* user */
     proc_info_creds(proc_i, task);

diff --git a/plugins/edriver-rust/src/bpf/common/print.h b/plugins/edriver-rust/src/bpf/common/print.h
@@ -21,7 +21,7 @@
         if (ctx->sbt == NULL)                                       \
             return 0;                                               \
         /* pass verifier */                                         \
-        if (ctx->offset > MAX_PERCPU_BUFSIZE - (MAX_STR + S_INT))   \
+        if (ctx->offset > MAX_PERCPU_BUFSIZE - MAX_STR - S_INT)     \
             return 0;                                               \
         int s = bpf_probe_read_str(                                 \
             &(ctx->sbt->buf[ctx->offset + S_INT]),                  \

diff --git a/plugins/edriver-rust/src/bpf/common/types.h b/plugins/edriver-rust/src/bpf/common/types.h
@@ -25,13 +25,19 @@ struct hds_cred {
 
 /* socket information */
 struct hds_socket_info {
-    __u16 family;
     __u32 local_address;
     __u16 local_port;
     __u32 remote_address;
     __u16 remote_port;
 };
 
+struct hds_socket_info_v6 {
+    struct in6_addr local_address;
+    __u16 local_port;
+    struct in6_addr remote_address;
+    __u16 remote_port;    
+};
+
 /* process information */
 struct proc_info {
     __u64 timestamp;
@@ -51,7 +57,9 @@ struct proc_info {
     char  ld_lib[MAX_STR_ENV];
     /* extra information */
     struct hds_cred cred;
+    __u16 family;
     struct hds_socket_info sinfo;
+    struct hds_socket_info_v6 sinfo_v6;
     char  pidtree[PIDTREE_LEN];
     /* others */
     __u16 pidtree_len;
@@ -68,5 +76,6 @@ struct hds_context {
 
 struct proc_info _proc SEC(".rodata") = {};
 struct hds_socket_info _sinfo SEC(".rodata") = {};
+struct hds_socket_info_v6 _sinfo_v6 SEC(".rodata") = {};
 
 #endif