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