xref: /linux/samples/bpf/sockex2_kern.c (revision bd628c1bed7902ec1f24ba0fe70758949146abbe)
1 #include <uapi/linux/bpf.h>
2 #include "bpf_helpers.h"
3 #include <uapi/linux/in.h>
4 #include <uapi/linux/if.h>
5 #include <uapi/linux/if_ether.h>
6 #include <uapi/linux/ip.h>
7 #include <uapi/linux/ipv6.h>
8 #include <uapi/linux/if_tunnel.h>
9 #define IP_MF		0x2000
10 #define IP_OFFSET	0x1FFF
11 
12 struct vlan_hdr {
13 	__be16 h_vlan_TCI;
14 	__be16 h_vlan_encapsulated_proto;
15 };
16 
17 struct flow_key_record {
18 	__be32 src;
19 	__be32 dst;
20 	union {
21 		__be32 ports;
22 		__be16 port16[2];
23 	};
24 	__u16 thoff;
25 	__u8 ip_proto;
26 };
27 
28 static inline int proto_ports_offset(__u64 proto)
29 {
30 	switch (proto) {
31 	case IPPROTO_TCP:
32 	case IPPROTO_UDP:
33 	case IPPROTO_DCCP:
34 	case IPPROTO_ESP:
35 	case IPPROTO_SCTP:
36 	case IPPROTO_UDPLITE:
37 		return 0;
38 	case IPPROTO_AH:
39 		return 4;
40 	default:
41 		return 0;
42 	}
43 }
44 
45 static inline int ip_is_fragment(struct __sk_buff *ctx, __u64 nhoff)
46 {
47 	return load_half(ctx, nhoff + offsetof(struct iphdr, frag_off))
48 		& (IP_MF | IP_OFFSET);
49 }
50 
51 static inline __u32 ipv6_addr_hash(struct __sk_buff *ctx, __u64 off)
52 {
53 	__u64 w0 = load_word(ctx, off);
54 	__u64 w1 = load_word(ctx, off + 4);
55 	__u64 w2 = load_word(ctx, off + 8);
56 	__u64 w3 = load_word(ctx, off + 12);
57 
58 	return (__u32)(w0 ^ w1 ^ w2 ^ w3);
59 }
60 
61 static inline __u64 parse_ip(struct __sk_buff *skb, __u64 nhoff, __u64 *ip_proto,
62 			     struct flow_key_record *flow)
63 {
64 	__u64 verlen;
65 
66 	if (unlikely(ip_is_fragment(skb, nhoff)))
67 		*ip_proto = 0;
68 	else
69 		*ip_proto = load_byte(skb, nhoff + offsetof(struct iphdr, protocol));
70 
71 	if (*ip_proto != IPPROTO_GRE) {
72 		flow->src = load_word(skb, nhoff + offsetof(struct iphdr, saddr));
73 		flow->dst = load_word(skb, nhoff + offsetof(struct iphdr, daddr));
74 	}
75 
76 	verlen = load_byte(skb, nhoff + 0/*offsetof(struct iphdr, ihl)*/);
77 	if (likely(verlen == 0x45))
78 		nhoff += 20;
79 	else
80 		nhoff += (verlen & 0xF) << 2;
81 
82 	return nhoff;
83 }
84 
85 static inline __u64 parse_ipv6(struct __sk_buff *skb, __u64 nhoff, __u64 *ip_proto,
86 			       struct flow_key_record *flow)
87 {
88 	*ip_proto = load_byte(skb,
89 			      nhoff + offsetof(struct ipv6hdr, nexthdr));
90 	flow->src = ipv6_addr_hash(skb,
91 				   nhoff + offsetof(struct ipv6hdr, saddr));
92 	flow->dst = ipv6_addr_hash(skb,
93 				   nhoff + offsetof(struct ipv6hdr, daddr));
94 	nhoff += sizeof(struct ipv6hdr);
95 
96 	return nhoff;
97 }
98 
99 static inline bool flow_dissector(struct __sk_buff *skb,
100 				  struct flow_key_record *flow)
101 {
102 	__u64 nhoff = ETH_HLEN;
103 	__u64 ip_proto;
104 	__u64 proto = load_half(skb, 12);
105 	int poff;
106 
107 	if (proto == ETH_P_8021AD) {
108 		proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
109 							h_vlan_encapsulated_proto));
110 		nhoff += sizeof(struct vlan_hdr);
111 	}
112 
113 	if (proto == ETH_P_8021Q) {
114 		proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
115 							h_vlan_encapsulated_proto));
116 		nhoff += sizeof(struct vlan_hdr);
117 	}
118 
119 	if (likely(proto == ETH_P_IP))
120 		nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
121 	else if (proto == ETH_P_IPV6)
122 		nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
123 	else
124 		return false;
125 
126 	switch (ip_proto) {
127 	case IPPROTO_GRE: {
128 		struct gre_hdr {
129 			__be16 flags;
130 			__be16 proto;
131 		};
132 
133 		__u64 gre_flags = load_half(skb,
134 					    nhoff + offsetof(struct gre_hdr, flags));
135 		__u64 gre_proto = load_half(skb,
136 					    nhoff + offsetof(struct gre_hdr, proto));
137 
138 		if (gre_flags & (GRE_VERSION|GRE_ROUTING))
139 			break;
140 
141 		proto = gre_proto;
142 		nhoff += 4;
143 		if (gre_flags & GRE_CSUM)
144 			nhoff += 4;
145 		if (gre_flags & GRE_KEY)
146 			nhoff += 4;
147 		if (gre_flags & GRE_SEQ)
148 			nhoff += 4;
149 
150 		if (proto == ETH_P_8021Q) {
151 			proto = load_half(skb,
152 					  nhoff + offsetof(struct vlan_hdr,
153 							   h_vlan_encapsulated_proto));
154 			nhoff += sizeof(struct vlan_hdr);
155 		}
156 
157 		if (proto == ETH_P_IP)
158 			nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
159 		else if (proto == ETH_P_IPV6)
160 			nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
161 		else
162 			return false;
163 		break;
164 	}
165 	case IPPROTO_IPIP:
166 		nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
167 		break;
168 	case IPPROTO_IPV6:
169 		nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
170 		break;
171 	default:
172 		break;
173 	}
174 
175 	flow->ip_proto = ip_proto;
176 	poff = proto_ports_offset(ip_proto);
177 	if (poff >= 0) {
178 		nhoff += poff;
179 		flow->ports = load_word(skb, nhoff);
180 	}
181 
182 	flow->thoff = (__u16) nhoff;
183 
184 	return true;
185 }
186 
187 struct pair {
188 	long packets;
189 	long bytes;
190 };
191 
192 struct bpf_map_def SEC("maps") hash_map = {
193 	.type = BPF_MAP_TYPE_HASH,
194 	.key_size = sizeof(__be32),
195 	.value_size = sizeof(struct pair),
196 	.max_entries = 1024,
197 };
198 
199 SEC("socket2")
200 int bpf_prog2(struct __sk_buff *skb)
201 {
202 	struct flow_key_record flow = {};
203 	struct pair *value;
204 	u32 key;
205 
206 	if (!flow_dissector(skb, &flow))
207 		return 0;
208 
209 	key = flow.dst;
210 	value = bpf_map_lookup_elem(&hash_map, &key);
211 	if (value) {
212 		__sync_fetch_and_add(&value->packets, 1);
213 		__sync_fetch_and_add(&value->bytes, skb->len);
214 	} else {
215 		struct pair val = {1, skb->len};
216 
217 		bpf_map_update_elem(&hash_map, &key, &val, BPF_ANY);
218 	}
219 	return 0;
220 }
221 
222 char _license[] SEC("license") = "GPL";
223