@@ -49,6 +49,7 @@ typedef struct statkey_t {
4949 __u16 src_port ; // source port
5050 __u16 dst_port ; // destination port
5151 __u8 proto ; // transport protocol
52+ pid_t pid ; // process ID
5253} statkey ;
5354
5455// Map value struct with counters
@@ -243,6 +244,70 @@ static inline void process_eth(void *data, void *data_end, __u64 pkt_len) {
243244 }
244245}
245246
247+ /**
248+ * Process the Ethernet header and extract relevant information, including the
249+ * process ID, to populate the key.
250+ *
251+ * @param data pointer to the start of the Ethernet header
252+ * @param data_end pointer to the end of the packet data
253+ * @param pkt_len length of the packet
254+ * @param pid process ID to be included in the key
255+ *
256+ * @return none
257+ *
258+ * @throws none
259+ */
260+ static inline void process_eth_pid (void * data , void * data_end , __u64 pkt_len ,
261+ pid_t pid ) {
262+ struct ethhdr * eth = data ;
263+
264+ // validate Ethernet size
265+ if ((void * )eth + sizeof (* eth ) > data_end ) {
266+ return ;
267+ }
268+
269+ // initialize key
270+ statkey key ;
271+ __builtin_memset (& key , 0 , sizeof (key ));
272+
273+ // populate PID (may be 0)
274+ key .pid = pid ;
275+
276+ // process only IPv4 and IPv6
277+ switch (bpf_ntohs (eth -> h_proto )) {
278+ case ETH_P_IP : {
279+ struct iphdr * ip4 = (void * )eth + sizeof (* eth );
280+
281+ if (process_ip4 (ip4 , data_end , & key ) == NOK )
282+ return ;
283+ }
284+
285+ break ;
286+ case ETH_P_IPV6 : {
287+ struct ipv6hdr * ip6 = (void * )eth + sizeof (* eth );
288+
289+ if (process_ip6 (ip6 , data_end , & key ) == NOK )
290+ return ;
291+ }
292+
293+ break ;
294+ default :
295+ return ;
296+ }
297+
298+ // lookup value in hash
299+ statvalue * val = (statvalue * )bpf_map_lookup_elem (& pkt_count , & key );
300+ if (val ) {
301+ // atomic XADD, doesn't need bpf_spin_lock()
302+ __sync_fetch_and_add (& val -> packets , 1 );
303+ __sync_fetch_and_add (& val -> bytes , pkt_len );
304+ } else {
305+ statvalue initval = {.packets = 1 , .bytes = pkt_len };
306+
307+ bpf_map_update_elem (& pkt_count , & key , & initval , BPF_NOEXIST );
308+ }
309+ }
310+
246311/**
247312 * Process the packet for traffic control and take necessary actions.
248313 *
@@ -259,6 +324,24 @@ static inline void tc_process_packet(struct __sk_buff *skb) {
259324 process_eth (data , data_end , skb -> len );
260325}
261326
327+ /**
328+ * Process the packet for traffic control and take necessary actions, with
329+ * the process ID included in the key.
330+ *
331+ * @param skb pointer to the packet buffer
332+ * @param pid process ID to be included in the key
333+ *
334+ * @return none
335+ *
336+ * @throws none
337+ */
338+ static inline void tc_process_packet_pid (struct __sk_buff * skb , pid_t pid ) {
339+ void * data = (void * )(long )skb -> data ;
340+ void * data_end = (void * )(long )skb -> data_end ;
341+
342+ process_eth_pid (data , data_end , skb -> len , pid );
343+ }
344+
262345/**
263346 * Process the packet for XDP (eXpress Data Path) and take necessary actions.
264347 *
@@ -275,8 +358,14 @@ static inline void xdp_process_packet(struct xdp_md *xdp) {
275358 process_eth (data , data_end , data_end - data );
276359}
277360
278- /*
279- * Main eBPF XDP program
361+ /**
362+ * Process the packet for XDP and take necessary actions.
363+ *
364+ * @param xdp pointer to the XDP context
365+ *
366+ * @return XDP_PASS
367+ *
368+ * @throws none
280369 */
281370SEC ("xdp" )
282371int xdp_count_packets (struct xdp_md * xdp ) {
@@ -285,8 +374,14 @@ int xdp_count_packets(struct xdp_md *xdp) {
285374 return XDP_PASS ;
286375}
287376
288- /*
289- * Main eBPF TC program
377+ /**
378+ * Process the packet for TC (Traffic Control) and take necessary actions.
379+ *
380+ * @param skb pointer to the skb
381+ *
382+ * @return TC_ACT_UNSPEC
383+ *
384+ * @throws none
290385 */
291386SEC ("tc" )
292387int tc_count_packets (struct __sk_buff * skb ) {
@@ -295,4 +390,24 @@ int tc_count_packets(struct __sk_buff *skb) {
295390 return TC_ACT_UNSPEC ;
296391}
297392
393+ /**
394+ * Process the packet for TC (Traffic Control) and take necessary actions. This
395+ * variant of the function also gets the current PID and passes it to the
396+ * `tc_process_packet` function.
397+ *
398+ * @param skb pointer to the skb
399+ *
400+ * @return TC_ACT_UNSPEC
401+ *
402+ * @throws none
403+ */
404+ SEC ("tc" )
405+ int tc_count_packets_pid (struct __sk_buff * skb ) {
406+ // get PID where possible
407+ pid_t pid = bpf_get_current_pid_tgid () & 0xFFFFFFFF ;
408+ tc_process_packet_pid (skb , pid );
409+
410+ return TC_ACT_UNSPEC ;
411+ }
412+
298413char __license [] SEC ("license" ) = "Dual MIT/GPL" ;
0 commit comments