first pass at DNS packet fetching

GPL/Events/EbpfEventProto.h

@@ -11,6 +11,8 @@
 #define EBPF_EVENTPROBE_EBPFEVENTPROTO_H
 
 #define TASK_COMM_LEN 16
+#define MAX_DNS_PACKET 512
+#define MAX_NR_SEGS 8
 
 #ifndef __KERNEL__
 #include <stdint.h>
@@ -42,6 +44,7 @@ enum ebpf_event_type {
     EBPF_EVENT_PROCESS_LOAD_MODULE          = (1 << 19),
     EBPF_EVENT_NETWORK_UDP_SENDMSG          = (1 << 20),
     EBPF_EVENT_NETWORK_UDP_RECVMSG          = (1 << 21),
+    EBPF_EVENT_NETWORK_DNS_PKT              = (1 << 22),
 };
 
 struct ebpf_event_header {
@@ -382,6 +385,29 @@ struct ebpf_net_event {
     char comm[TASK_COMM_LEN];
 } __attribute__((packed));
 
+struct dns_pkt_header {
+    uint16_t transaction_id;
+    uint16_t flags;
+    uint16_t num_questions;
+    uint16_t num_answers;
+    uint16_t num_auth_rrs;
+    uint16_t num_additional_rrs;
+} __attribute__((packed));
+
+struct dns_body {
+    size_t len;
+    uint8_t pkt[MAX_DNS_PACKET];
+} __attribute((packed));
+
+struct ebpf_dns_event {
+    struct ebpf_event_header hdr;
+    struct ebpf_pid_info pids;
+    struct ebpf_net_info net;
+    char comm[TASK_COMM_LEN];
+    enum ebpf_event_type udp_evt;
+    struct dns_body pkts[MAX_NR_SEGS];
+} __attribute__((packed));
+
 // Basic event statistics
 struct ebpf_event_stats {
     uint64_t lost; // lost events due to a full ringbuffer

GPL/Events/Helpers.h

@@ -138,6 +138,9 @@ const volatile int consumer_pid = 0;
 // From include/uapi/asm-generic/termbits.h
 #define ECHO 0x00008
 
+/* tty_write */
+DECL_FIELD_OFFSET(iov_iter, __iov);
+
 static bool IS_ERR_OR_NULL(const void *ptr)
 {
     return (!ptr) || (unsigned long)ptr >= (unsigned long)-MAX_ERRNO;
@@ -357,4 +360,25 @@ static int is_equal_prefix(const char *str1, const char *str2, int len)
     return !strncmp(str1, str2, len);
 }
 
+static int get_iovec_nr_segs_or_max(struct iov_iter *from)
+{
+    u64 nr_segs = BPF_CORE_READ(from, nr_segs);
+    nr_segs     = nr_segs > MAX_NR_SEGS ? MAX_NR_SEGS : nr_segs;
+    return nr_segs;
+}
+
+struct udp_ctx {
+    struct sock *sk;
+    struct msghdr *hdr;
+    int flags;
+} __attribute__((packed));
+
+// scratchspace map for fetching the arguments from a kretprobe
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __type(key, u64);
+    __type(value, struct udp_ctx);
+    __uint(max_entries, 1024);
+} pkt_ctx SEC(".maps");
+
 #endif // EBPF_EVENTPROBE_HELPERS_H

GPL/Events/Network/Network.h

@@ -14,6 +14,14 @@
 #define AF_INET 2
 #define AF_INET6 10
 
+#define MSG_PEEK 2
+
+// See RFC1035
+#define DNS_QR_BIT 1 << 15
+
+// I have made this number up to make the verifier happy
+#define DNS_MAX_LABELS 255
+
 static int ebpf_sock_info__fill(struct ebpf_net_info *net, struct sock *sk)
 {
     int err = 0;

GPL/Events/Network/Probe.bpf.c

@@ -20,6 +20,107 @@
 
 DECL_FUNC_RET(inet_csk_accept);
 
+static int sock_dns_event_handle(struct sock *sk,
+                                 struct msghdr *msg,
+                                 enum ebpf_event_type evt_type,
+                                 size_t size)
+{
+    if (!sk) {
+        return 0;
+    }
+
+    if (ebpf_events_is_trusted_pid())
+        return 0;
+
+    struct ebpf_dns_event *event = bpf_ringbuf_reserve(&ringbuf, sizeof(*event), 0);
+    if (!event)
+        return 0;
+
+    // fill in socket and process metadata
+    if (ebpf_sock_info__fill(&event->net, sk)) {
+        goto out;
+    }
+
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+    ebpf_pid_info__fill(&event->pids, task);
+    bpf_get_current_comm(event->comm, TASK_COMM_LEN);
+    event->hdr.ts = bpf_ktime_get_ns();
+
+    // filter out non-dns packets
+    if (event->net.dport != 53 && event->net.sport != 53) {
+        goto out;
+    }
+
+    // deal with the iovec_iter type
+    // newer kernels added a ubuf type to the iov_iter union,
+    // which post-dates our vmlinux, but also they added ITER_UBUF as the
+    // first value in the iter_type enum, which makes checking it a tad hard.
+    // In theory we should be able to read from both types as long as we're careful
+
+    struct iov_iter *from = &msg->msg_iter;
+
+    u64 nr_segs    = get_iovec_nr_segs_or_max(from);
+    u64 iovec_size = BPF_CORE_READ(from, count);
+
+    const struct iovec *iov;
+    if (FIELD_OFFSET(iov_iter, __iov))
+        iov = (const struct iovec *)((char *)from + FIELD_OFFSET(iov_iter, __iov));
+    else if (bpf_core_field_exists(from->iov))
+        iov = BPF_CORE_READ(from, iov);
+    else {
+        bpf_printk("unknown offset in iovec structure, bug?");
+        goto out;
+    }
+
+    if (nr_segs == 1) {
+        // actually read in raw packet data
+        // use the retvalue of recvmsg/the count value of sendmsg instead of the the iovec count
+        // the count of the iovec in udp_recvmsg is the size of the buffer, not the size of the
+        // bytes read.
+        void *base         = BPF_CORE_READ(iov, iov_base);
+        event->pkts[0].len = size;
+        // make verifier happy, we can't have an out-of-bounds write
+        if (size > MAX_DNS_PACKET) {
+            bpf_printk("size of packet (%d) exceeds max packet size (%d), skipping", size,
+                       MAX_DNS_PACKET);
+            goto out;
+        }
+        long readok = bpf_probe_read(event->pkts[0].pkt, size, base);
+        if (readok != 0) {
+            bpf_printk("invalid read from iovec structure: %d", readok);
+            goto out;
+        }
+    } else {
+        // we have multiple segments.
+        // Can't rely on the size value from the function, revert to the iovec size to read into the
+        // buffer
+        // In practice, I haven't seen a DNS packet with more than one iovec segment;
+        // the size of UDP DNS packet is limited to 512 bytes, so not sure if this is possible?
+        for (int seg = 0; seg < nr_segs; seg++) {
+            if (seg >= MAX_NR_SEGS)
+                goto out;
+
+            struct iovec *cur_iov = (struct iovec *)&iov[seg];
+            void *base            = BPF_CORE_READ(cur_iov, iov_base);
+            size_t bufsize        = BPF_CORE_READ(cur_iov, iov_len);
+            event->pkts[seg].len  = bufsize;
+            if (bufsize > sizeof(event->pkts[seg].pkt)) {
+                goto out;
+            }
+            bpf_probe_read(event->pkts[seg].pkt, bufsize, base);
+        }
+    }
+
+    event->hdr.type = EBPF_EVENT_NETWORK_DNS_PKT;
+    event->udp_evt  = evt_type;
+    bpf_ringbuf_submit(event, 0);
+    return 0;
+
+out:
+    bpf_ringbuf_discard(event, 0);
+    return 0;
+}
+
 static int sock_object_handle(struct sock *sk, enum ebpf_event_type evt_type)
 {
     if (!sk)
@@ -43,31 +144,100 @@ static int sock_object_handle(struct sock *sk, enum ebpf_event_type evt_type)
     return 0;
 }
 
+/*
+=============================== DNS probes ===============================
+*/
+
 SEC("fentry/udp_sendmsg")
-int BPF_PROG(fentry__udp_sendmsg, struct sock *sk)
+int BPF_PROG(fentry__udp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size)
 {
-    bpf_printk("got ip4_datagram event");
-    return sock_object_handle(sk, EBPF_EVENT_NETWORK_UDP_SENDMSG);
+    return sock_dns_event_handle(sk, msg, EBPF_EVENT_NETWORK_UDP_SENDMSG, size);
 }
 
 SEC("fexit/udp_recvmsg")
-int BPF_PROG(fexit__udp_recvmsg, struct sock *sk)
+int BPF_PROG(fexit__udp_recvmsg,
+             struct sock *sk,
+             struct msghdr *msg,
+             size_t len,
+             int flags,
+             int *addr_len,
+             int ret)
 {
-    return sock_object_handle(sk, EBPF_EVENT_NETWORK_UDP_RECVMSG);
+    // check the peeking flag; if set to peek, the msghdr won't contain any data
+    if (flags & MSG_PEEK) {
+        return 0;
+    }
+    return sock_dns_event_handle(sk, msg, EBPF_EVENT_NETWORK_UDP_RECVMSG, ret);
 }
 
 SEC("kprobe/udp_sendmsg")
-int BPF_KPROBE(kprobe__udp_sendmsg, struct sock *sk)
+int BPF_KPROBE(kprobe__udp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size)
 {
-    return sock_object_handle(sk, EBPF_EVENT_NETWORK_UDP_SENDMSG);
+    return sock_dns_event_handle(sk, msg, EBPF_EVENT_NETWORK_UDP_SENDMSG, size);
 }
 
+// We can't get the arguments from a kretprobe, so instead save off the pointer in
+// in the kprobe, then fetch the pointer from a context map in the kretprobe
+
 SEC("kprobe/udp_recvmsg")
-int BPF_KPROBE(kprobe__udp_recvmsg, struct sock *sk)
+int BPF_KPROBE(
+    kprobe__udp_recvmsg, struct sock *sk, struct msghdr *msg, size_t len, int flags, int *addr_len)
 {
-    return sock_object_handle(sk, EBPF_EVENT_NETWORK_UDP_RECVMSG);
+    struct udp_ctx kctx;
+    kctx.flags = flags;
+
+    // I suspect that using the PID_TID isn't the most reliable way to map the sockets/iters
+    // not sure what else we could use that's accessable from the kretprobe, though.
+    u64 pid_tid = bpf_get_current_pid_tgid();
+
+    long iter_err = bpf_probe_read(&kctx.hdr, sizeof(kctx.hdr), &msg);
+    if (iter_err != 0) {
+        bpf_printk("error reading msg_iter in udp_recvmsg: %d", iter_err);
+        return 0;
+    }
+
+    long sk_err = bpf_probe_read(&kctx.sk, sizeof(kctx.sk), &sk);
+    if (sk_err != 0) {
+        bpf_printk("error reading msg_iter in udp_recvmsg: %d", sk_err);
+        return 0;
+    }
+
+    long update_err = bpf_map_update_elem(&pkt_ctx, &pid_tid, &kctx, BPF_ANY);
+    if (update_err != 0) {
+        bpf_printk("error updating context map in udp_recvmsg: %d", update_err);
+        return 0;
+    }
+
+    return 0;
 }
 
+SEC("kretprobe/udp_recvmsg")
+int BPF_KRETPROBE(kretprobe__udp_recvmsg, int ret)
+{
+    bpf_printk("in kretprobe udp_recvmsg....");
+
+    u64 pid_tid = bpf_get_current_pid_tgid();
+
+    void *vctx = bpf_map_lookup_elem(&pkt_ctx, &pid_tid);
+
+    struct udp_ctx kctx;
+    long read_err = bpf_probe_read(&kctx, sizeof(kctx), vctx);
+    if (read_err != 0) {
+        bpf_printk("error reading back context in udp_recvmsg: %d", read_err);
+    }
+
+    // check the peeking flag; if set to peek, the msghdr won't contain any data
+    if (kctx.flags & MSG_PEEK) {
+        return 0;
+    }
+
+    return sock_dns_event_handle(kctx.sk, kctx.hdr, EBPF_EVENT_NETWORK_UDP_RECVMSG, ret);
+}
+
+/*
+=============================== TCP probes ===============================
+*/
+
 SEC("fexit/inet_csk_accept")
 int BPF_PROG(fexit__inet_csk_accept)
 {

GPL/Events/Process/Probe.bpf.c

@@ -18,9 +18,6 @@
 #include "State.h"
 #include "Varlen.h"
 
-/* tty_write */
-DECL_FIELD_OFFSET(iov_iter, __iov);
-
 // Limits on large things we send up as variable length parameters.
 //
 // These should be kept _well_ under half the size of the event_buffer_map or
@@ -526,8 +523,6 @@ int BPF_KPROBE(kprobe__commit_creds, struct cred *new)
     return commit_creds__enter(new);
 }
 
-#define MAX_NR_SEGS 8
-
 static int output_tty_event(struct ebpf_tty_dev *slave, const void *base, size_t base_len)
 {
     struct ebpf_process_tty_write_event *event;
@@ -611,8 +606,7 @@ static int tty_write__enter(struct kiocb *iocb, struct iov_iter *from)
     else
         goto out;
 
-    u64 nr_segs = BPF_CORE_READ(from, nr_segs);
-    nr_segs     = nr_segs > MAX_NR_SEGS ? MAX_NR_SEGS : nr_segs;
+    u64 nr_segs = get_iovec_nr_segs_or_max(from);
 
     if (nr_segs == 0) {
         u64 count = BPF_CORE_READ(from, count);