Split iovec reading for tcp_recvmsg

bpf/http_sock.c

@@ -29,7 +29,7 @@ struct {
 // Temporary tracking of tcp_recvmsg arguments
 typedef struct recv_args {
     u64 sock_ptr; // linux sock or socket address
-    u64 iovec_ptr;
+    struct iov_iter_data iovec_ctx;
 } recv_args_t;
 
 struct {
@@ -426,9 +426,10 @@ int BPF_KPROBE(kprobe_tcp_recvmsg, struct sock *sk, struct msghdr *msg, size_t l
 
     recv_args_t args = {
         .sock_ptr = (u64)sk,
-        .iovec_ptr = (u64)(msg)
     };
 
+    get_iovec_ctx((iovec_iter_ctx *)&args.iovec_ctx, msg);
+
     bpf_map_update_elem(&active_recv_args, &id, &args, BPF_ANY);
 
     return 0;
@@ -443,7 +444,9 @@ static __always_inline int return_recvmsg(void *ctx, u64 id, int copied_len) {
         goto done;
     }
 
-    if (!args->iovec_ptr) {
+    iovec_iter_ctx *iov_ctx = (iovec_iter_ctx *)&args->iovec_ctx;
+
+    if (!iov_ctx->iov && !iov_ctx->ubuf) {
         bpf_dbg_printk("iovec_ptr found in kprobe is NULL, ignoring this tcp_recvmsg");
         bpf_map_delete_elem(&active_recv_args, &id);
 
@@ -452,7 +455,6 @@ static __always_inline int return_recvmsg(void *ctx, u64 id, int copied_len) {
 
     pid_connection_info_t info = {};
 
-    void *iovec_ptr = (void *)args->iovec_ptr;
     void *sock_ptr = (void *)args->sock_ptr;
 
     bpf_map_delete_elem(&active_recv_args, &id);
@@ -470,7 +472,7 @@ static __always_inline int return_recvmsg(void *ctx, u64 id, int copied_len) {
             if (!active_ssl) {
                 u8* buf = iovec_memory();
                 if (buf) {
-                    copied_len = read_msghdr_buf((void *)iovec_ptr, buf, copied_len);
+                    copied_len = read_iovec_ctx(iov_ctx, buf, copied_len);
                     if (copied_len) {
                         // doesn't return must be logically last statement
                         handle_buf_with_connection(ctx, &info, buf, copied_len, NO_SSL, TCP_RECV, orig_dport);

bpf/protocol_common.h

@@ -93,6 +93,17 @@ static __always_inline http_connection_metadata_t *connection_meta_by_direction(
     return meta;
 }
 
+// This data structure must be kept in sync with iov_iter__dummy below. All pointers
+// must be converted to u64 to have enough data for 64bit (and 32bit) platforms.
+struct iov_iter_data {
+    unsigned int __dummy_type;
+    u8 __dummy_iter_type;
+    u64 __dummy_ubuf;
+    u64 __dummy_iov;
+    u64 __dummy_iov1;
+    unsigned long __dummy_nr_segs;
+};
+
 struct iov_iter___dummy {
     unsigned int type; // for co-re support, use iter_type instead
     u8 iter_type;
@@ -110,6 +121,8 @@ enum iter_type___dummy {
 
 // extracts kernel specific iov_iter information into a iovec_iter_ctx instance
 static __always_inline void get_iovec_ctx(iovec_iter_ctx* ctx, struct msghdr *msg) {
+    ctx->ubuf = NULL;
+    ctx->iov = NULL;
     if (bpf_core_field_exists(((struct iov_iter___dummy*)&msg->msg_iter)->type)) {
         // clear the direction bit when reading iovec_iter::type to end up
         // with the original enumerator value (the direction bit is the LSB
@@ -136,53 +149,49 @@ static __always_inline void get_iovec_ctx(iovec_iter_ctx* ctx, struct msghdr *ms
     ctx->nr_segs = BPF_CORE_READ((struct iov_iter___dummy*)&msg->msg_iter, nr_segs);
 }
 
-static __always_inline int read_msghdr_buf(struct msghdr *msg, u8* buf, size_t max_len) {
+static __always_inline int read_iovec_ctx(iovec_iter_ctx *ctx, u8* buf, size_t max_len) {
     if (max_len == 0) {
         return 0;
     }
 
     bpf_clamp_umax(max_len, IO_VEC_MAX_LEN);
 
-    iovec_iter_ctx ctx;
-
-    get_iovec_ctx(&ctx, msg);
-
-    bpf_dbg_printk("iter_type=%u", ctx.iter_type);
-    bpf_dbg_printk("nr_segs=%lu, iov=%p, ubuf=%p", ctx.nr_segs, ctx.iov, ctx.ubuf);
+    bpf_dbg_printk("iter_type=%u", ctx->iter_type);
+    bpf_dbg_printk("nr_segs=%lu, iov=%p, ubuf=%p", ctx->nr_segs, ctx->iov, ctx->ubuf);
 
     // ITER_UBUF only exists in kernels >= 6.0 - earlier kernels use ITER_IOVEC
     if (bpf_core_enum_value_exists(enum iter_type___dummy, ITER_UBUF)) {
         const int iter_ubuf = bpf_core_enum_value(enum iter_type___dummy, ITER_UBUF);
 
         // ITER_UBUF is never a bitmask, and can be 0, so we perform a proper
         // equality check rather than a bitwise and like we do for ITER_IOVEC
-        if (ctx.ubuf != NULL && ctx.iter_type == iter_ubuf) {
+        if (ctx->ubuf != NULL && ctx->iter_type == iter_ubuf) {
             bpf_clamp_umax(max_len, IO_VEC_MAX_LEN);
-            return bpf_probe_read(buf, max_len, ctx.ubuf) == 0 ? max_len : 0;
+            return bpf_probe_read(buf, max_len, ctx->ubuf) == 0 ? max_len : 0;
         }
     }
 
     const int iter_iovec = bpf_core_enum_value(enum iter_type, ITER_IOVEC);
 
-    if (ctx.iter_type != iter_iovec) {
+    if (ctx->iter_type != iter_iovec) {
         return 0;
     }
 
     u32 tot_len = 0;
 
     enum { max_segments = 4 };
 
-    bpf_clamp_umax(ctx.nr_segs, max_segments);
+    bpf_clamp_umax(ctx->nr_segs, max_segments);
 
     // Loop couple of times reading the various io_vecs
-    for (unsigned long i = 0; i < ctx.nr_segs && i < max_segments; i++) {
+    for (unsigned long i = 0; i < ctx->nr_segs && i < max_segments; i++) {
         struct iovec vec;
 
-        if (bpf_probe_read_kernel(&vec, sizeof(vec), &ctx.iov[i]) != 0) {
+        if (bpf_probe_read_kernel(&vec, sizeof(vec), &ctx->iov[i]) != 0) {
             break;
         }
 
-        // bpf_dbg_printk("iov[%d]=%llx", i, &ctx.iov[i]);
+        // bpf_dbg_printk("iov[%d]=%llx", i, &ctx->iov[i]);
         // bpf_dbg_printk("base %llx, len %d", vec.iov_base, vec.iov_len);
 
         if (!vec.iov_base || !vec.iov_len) {
@@ -211,6 +220,18 @@ static __always_inline int read_msghdr_buf(struct msghdr *msg, u8* buf, size_t m
     return tot_len;
 }
 
+static __always_inline int read_msghdr_buf(struct msghdr *msg, u8* buf, size_t max_len) {
+    if (max_len == 0) {
+        return 0;
+    }
+
+    iovec_iter_ctx ctx;
+
+    get_iovec_ctx(&ctx, msg);
+
+    return read_iovec_ctx(&ctx, buf, max_len);
+}
+
 // We sort the connection info to ensure we can track requests and responses. However, if the destination port
 // is somehow in the ephemeral port range, it can be higher than the source port and we'd use the sorted connection
 // info in user space, effectively reversing the flow of the operation. We keep track of the original destination port