1 // Copyright (c) PLUMgrid, Inc.
2 // Licensed under the Apache License, Version 2.0 (the "License")
3 
4 #include <bcc/proto.h>
5 
6 struct IPKey {
7   u32 dip;
8   u32 sip;
9 };
10 struct IPLeaf {
11   u64 rx_pkts;
12   u64 tx_pkts;
13 };
14 
15 BPF_HASH(stats, struct IPKey, struct IPLeaf, 256);
16 
on_packet(struct __sk_buff * skb)17 int on_packet(struct __sk_buff *skb) {
18   u8 *cursor = 0;
19   ethernet: {
20     struct ethernet_t *ethernet = cursor_advance(cursor, sizeof(*ethernet));
21     switch (ethernet->type) {
22         case ETH_P_IP: goto ip;
23         case ETH_P_8021Q: goto dot1q;
24         default: goto EOP;
25     }
26   }
27 
28   dot1q: {
29     struct dot1q_t *dot1q = cursor_advance(cursor, sizeof(*dot1q));
30     switch (dot1q->type) {
31       case ETH_P_8021Q: goto ip;
32       default: goto EOP;
33     }
34   }
35 
36   ip: {
37     struct ip_t *ip = cursor_advance(cursor, sizeof(*ip));
38     int rx = 0, tx = 0;
39     struct IPKey key;
40     if (ip->dst > ip->src) {
41       key.dip = ip->dst;
42       key.sip = ip->src;
43       rx = 1;
44     } else {
45       key.dip = ip->src;
46       key.sip = ip->dst;
47       tx = 1;
48     }
49     struct IPLeaf zleaf = {0};
50     struct IPLeaf *leaf = stats.lookup_or_init(&key, &zleaf);
51     lock_xadd(&leaf->rx_pkts, rx);
52     lock_xadd(&leaf->tx_pkts, tx);
53   }
54 
55 EOP:
56   return 0;
57 }
58