[CWS] Fill span context of network events (#38468)

Author: YoannGh

pkg/security/ebpf/c/include/helpers/network/dns.h

@@ -30,6 +30,14 @@ __attribute__((always_inline)) struct dns_event_t *reset_dns_event(struct __sk_b
     // process context
     fill_network_process_context_from_pkt(&evt->process, pkt);
 
+    u64 sched_cls_has_current_pid_tgid_helper = 0;
+    LOAD_CONSTANT("sched_cls_has_current_pid_tgid_helper", sched_cls_has_current_pid_tgid_helper);
+    if (sched_cls_has_current_pid_tgid_helper) {
+        // reset and fill span context
+        reset_span_context(&evt->span);
+        fill_span_context(&evt->span);
+    }
+
     // network context
     fill_network_context(&evt->network, skb, pkt);
 

pkg/security/ebpf/c/include/helpers/network/imds.h

@@ -26,6 +26,14 @@ __attribute__((always_inline)) struct imds_event_t *reset_imds_event(struct __sk
     // process context
     fill_network_process_context_from_pkt(&evt->process, pkt);
 
+    u64 sched_cls_has_current_pid_tgid_helper = 0;
+    LOAD_CONSTANT("sched_cls_has_current_pid_tgid_helper", sched_cls_has_current_pid_tgid_helper);
+    if (sched_cls_has_current_pid_tgid_helper) {
+        // reset and fill span context
+        reset_span_context(&evt->span);
+        fill_span_context(&evt->span);
+    }
+
     // network context
     fill_network_context(&evt->network, skb, pkt);
 

pkg/security/ebpf/c/include/helpers/network/stats.h

@@ -49,6 +49,14 @@ __attribute__((always_inline)) int flush_network_stats(u32 pid, struct active_fl
     // process context
     fill_network_process_context(&evt->process, pid, entry->netns);
 
+    u64 sched_cls_has_current_pid_tgid_helper = 0;
+    LOAD_CONSTANT("sched_cls_has_current_pid_tgid_helper", sched_cls_has_current_pid_tgid_helper);
+    if (sched_cls_has_current_pid_tgid_helper) {
+        // reset and fill span context
+        reset_span_context(&evt->span);
+        fill_span_context(&evt->span);
+    }
+
     // network context
     fill_network_device_context(&evt->device, entry->netns, entry->ifindex);
 

pkg/security/ebpf/c/include/helpers/span.h

@@ -49,6 +49,12 @@ void __attribute__((always_inline)) fill_span_context(struct span_context_t *spa
     }
 }
 
+void __attribute__((always_inline)) reset_span_context(struct span_context_t *span) {
+    span->span_id = 0;
+    span->trace_id[0] = 0;
+    span->trace_id[1] = 0;
+}
+
 void __attribute__((always_inline)) copy_span_context(struct span_context_t *src, struct span_context_t *dst) {
     dst->span_id = src->span_id;
     dst->trace_id[0] = src->trace_id[0];

pkg/security/ebpf/c/include/hooks/network/dns.h

@@ -186,6 +186,12 @@ TAIL_CALL_CLASSIFIER_FNC(dns_response, struct __sk_buff *skb) {
     } else {
         send_packet_with_context = true;
         fill_network_process_context_from_pkt(&map_elem->full_dns_response.process, pkt);
+        u64 sched_cls_has_current_pid_tgid_helper = 0;
+        LOAD_CONSTANT("sched_cls_has_current_pid_tgid_helper", sched_cls_has_current_pid_tgid_helper);
+        if (sched_cls_has_current_pid_tgid_helper) {
+            // fill span context (that was previously reset by reset_dns_response_event)
+            fill_span_context(&map_elem->full_dns_response.span);
+        }
         fill_network_context(&map_elem->full_dns_response.network, skb, pkt);
         err = bpf_skb_load_bytes(skb, pkt->offset, &map_elem->full_dns_response.header, sizeof(struct dnshdr));
         header_id = map_elem->full_dns_response.header.id;

pkg/security/ebpf/c/include/hooks/network/raw.h

@@ -28,6 +28,14 @@ TAIL_CALL_CLASSIFIER_FNC(raw_packet_sender, struct __sk_buff *skb) {
     // process context
     fill_network_process_context_from_pkt(&evt->process, pkt);
 
+    u64 sched_cls_has_current_pid_tgid_helper = 0;
+    LOAD_CONSTANT("sched_cls_has_current_pid_tgid_helper", sched_cls_has_current_pid_tgid_helper);
+    if (sched_cls_has_current_pid_tgid_helper) {
+        // reset and fill span context
+        reset_span_context(&evt->span);
+        fill_span_context(&evt->span);
+    }
+
     struct proc_cache_t *entry = get_proc_cache(evt->process.pid);
     if (entry == NULL) {
         evt->container.container_id[0] = 0;

pkg/security/ebpf/kernel/kernel_bpf.go

@@ -170,3 +170,9 @@ func (k *Version) SupportCORE() bool {
 	_, err := btf.LoadKernelSpec()
 	return err == nil
 }
+
+// HasBpfGetCurrentPidTgidForSchedCLS returns true if the kernel supports bpf_get_current_pid_tgid for Sched CLS program type
+// https://github.com/torvalds/linux/commit/eb166e522c77699fc19bfa705652327a1e51a117
+func (k *Version) HasBpfGetCurrentPidTgidForSchedCLS() bool {
+	return features.HaveProgramHelper(ebpf.SchedCLS, asm.FnGetCurrentPidTgid) == nil
+}

pkg/security/ebpf/kernel/kernel_nobpf.go

@@ -73,3 +73,9 @@ func (k *Version) HasSKStorageInTracingPrograms() bool {
 func (k *Version) HasTracingHelpersInCgroupSysctlPrograms() bool {
 	return false
 }
+
+// HasBpfGetCurrentPidTgidForSchedCLS returns true if the kernel supports bpf_get_current_pid_tgid for Sched CLS program type
+// https://github.com/torvalds/linux/commit/eb166e522c77699fc19bfa705652327a1e51a117
+func (k *Version) HasBpfGetCurrentPidTgidForSchedCLS() bool {
+	return false
+}

pkg/security/probe/probe_ebpf.go

@@ -2332,6 +2332,10 @@ func (p *EBPFProbe) initManagerOptionsConstants() {
 			Name:  "raw_packet_filter",
 			Value: utils.BoolTouint64(p.config.Probe.NetworkRawPacketFilter != "none"),
 		},
+		manager.ConstantEditor{
+			Name:  "sched_cls_has_current_pid_tgid_helper",
+			Value: utils.BoolTouint64(p.kernelVersion.HasBpfGetCurrentPidTgidForSchedCLS()),
+		},
 	)
 
 	if p.kernelVersion.HavePIDLinkStruct() {