[mirror_ubuntu-artful-kernel.git] / samples / bpf / sockex2_kern.c

#include <uapi/linux/bpf.h>
#include "bpf_helpers.h"
#include <uapi/linux/in.h>
#include <uapi/linux/if.h>
#include <uapi/linux/if_ether.h>
#include <uapi/linux/ip.h>
#include <uapi/linux/ipv6.h>
#include <uapi/linux/if_tunnel.h>
#define IP_MF		0x2000
#define IP_OFFSET	0x1FFF

struct vlan_hdr {
	__be16 h_vlan_TCI;
	__be16 h_vlan_encapsulated_proto;
};

struct flow_keys {
	__be32 src;
	__be32 dst;
	union {
		__be32 ports;
		__be16 port16[2];
	};
	__u16 thoff;
	__u8 ip_proto;
};

static inline int proto_ports_offset(__u64 proto)
{
	switch (proto) {
	case IPPROTO_TCP:
	case IPPROTO_UDP:
	case IPPROTO_DCCP:
	case IPPROTO_ESP:
	case IPPROTO_SCTP:
	case IPPROTO_UDPLITE:
		return 0;
	case IPPROTO_AH:
		return 4;
	default:
		return 0;
	}
}

static inline int ip_is_fragment(struct __sk_buff *ctx, __u64 nhoff)
{
	return load_half(ctx, nhoff + offsetof(struct iphdr, frag_off))
		& (IP_MF | IP_OFFSET);
}

static inline __u32 ipv6_addr_hash(struct __sk_buff *ctx, __u64 off)
{
	__u64 w0 = load_word(ctx, off);
	__u64 w1 = load_word(ctx, off + 4);
	__u64 w2 = load_word(ctx, off + 8);
	__u64 w3 = load_word(ctx, off + 12);

	return (__u32)(w0 ^ w1 ^ w2 ^ w3);
}

static inline __u64 parse_ip(struct __sk_buff *skb, __u64 nhoff, __u64 *ip_proto,
			     struct flow_keys *flow)
{
	__u64 verlen;

	if (unlikely(ip_is_fragment(skb, nhoff)))
		*ip_proto = 0;
	else
		*ip_proto = load_byte(skb, nhoff + offsetof(struct iphdr, protocol));

	if (*ip_proto != IPPROTO_GRE) {
		flow->src = load_word(skb, nhoff + offsetof(struct iphdr, saddr));
		flow->dst = load_word(skb, nhoff + offsetof(struct iphdr, daddr));
	}

	verlen = load_byte(skb, nhoff + 0/*offsetof(struct iphdr, ihl)*/);
	if (likely(verlen == 0x45))
		nhoff += 20;
	else
		nhoff += (verlen & 0xF) << 2;

	return nhoff;
}

static inline __u64 parse_ipv6(struct __sk_buff *skb, __u64 nhoff, __u64 *ip_proto,
			       struct flow_keys *flow)
{
	*ip_proto = load_byte(skb,
			      nhoff + offsetof(struct ipv6hdr, nexthdr));
	flow->src = ipv6_addr_hash(skb,
				   nhoff + offsetof(struct ipv6hdr, saddr));
	flow->dst = ipv6_addr_hash(skb,
				   nhoff + offsetof(struct ipv6hdr, daddr));
	nhoff += sizeof(struct ipv6hdr);

	return nhoff;
}

static inline bool flow_dissector(struct __sk_buff *skb, struct flow_keys *flow)
{
	__u64 nhoff = ETH_HLEN;
	__u64 ip_proto;
	__u64 proto = load_half(skb, 12);
	int poff;

	if (proto == ETH_P_8021AD) {
		proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
							h_vlan_encapsulated_proto));
		nhoff += sizeof(struct vlan_hdr);
	}

	if (proto == ETH_P_8021Q) {
		proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
							h_vlan_encapsulated_proto));
		nhoff += sizeof(struct vlan_hdr);
	}

	if (likely(proto == ETH_P_IP))
		nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
	else if (proto == ETH_P_IPV6)
		nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
	else
		return false;

	switch (ip_proto) {
	case IPPROTO_GRE: {
		struct gre_hdr {
			__be16 flags;
			__be16 proto;
		};

		__u64 gre_flags = load_half(skb,
					    nhoff + offsetof(struct gre_hdr, flags));
		__u64 gre_proto = load_half(skb,
					    nhoff + offsetof(struct gre_hdr, proto));

		if (gre_flags & (GRE_VERSION|GRE_ROUTING))
			break;

		proto = gre_proto;
		nhoff += 4;
		if (gre_flags & GRE_CSUM)
			nhoff += 4;
		if (gre_flags & GRE_KEY)
			nhoff += 4;
		if (gre_flags & GRE_SEQ)
			nhoff += 4;

		if (proto == ETH_P_8021Q) {
			proto = load_half(skb,
					  nhoff + offsetof(struct vlan_hdr,
							   h_vlan_encapsulated_proto));
			nhoff += sizeof(struct vlan_hdr);
		}

		if (proto == ETH_P_IP)
			nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
		else if (proto == ETH_P_IPV6)
			nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
		else
			return false;
		break;
	}
	case IPPROTO_IPIP:
		nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
		break;
	case IPPROTO_IPV6:
		nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
		break;
	default:
		break;
	}

	flow->ip_proto = ip_proto;
	poff = proto_ports_offset(ip_proto);
	if (poff >= 0) {
		nhoff += poff;
		flow->ports = load_word(skb, nhoff);
	}

	flow->thoff = (__u16) nhoff;

	return true;
}

struct pair {
	long packets;
	long bytes;
};

struct bpf_map_def SEC("maps") hash_map = {
	.type = BPF_MAP_TYPE_HASH,
	.key_size = sizeof(__be32),
	.value_size = sizeof(struct pair),
	.max_entries = 1024,
};

SEC("socket2")
int bpf_prog2(struct __sk_buff *skb)
{
	struct flow_keys flow;
	struct pair *value;
	u32 key;

	if (!flow_dissector(skb, &flow))
		return 0;

	key = flow.dst;
	value = bpf_map_lookup_elem(&hash_map, &key);
	if (value) {
		__sync_fetch_and_add(&value->packets, 1);
		__sync_fetch_and_add(&value->bytes, skb->len);
	} else {
		struct pair val = {1, skb->len};

		bpf_map_update_elem(&hash_map, &key, &val, BPF_ANY);
	}
	return 0;
}

char _license[] SEC("license") = "GPL";
Commit	Line	Data
fbe33108 AS	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_keys {
	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
614cd3bd	45	static inline int ip_is_fragment(struct __sk_buff *ctx, __u64 nhoff)
fbe33108 AS	46	{
	47	return load_half(ctx, nhoff + offsetof(struct iphdr, frag_off))
	48	& (IP_MF \| IP_OFFSET);
	49	}
	50
614cd3bd	51	static inline __u32 ipv6_addr_hash(struct __sk_buff *ctx, __u64 off)
fbe33108 AS	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
614cd3bd	61	static inline __u64 parse_ip(struct __sk_buff skb, __u64 nhoff, __u64 ip_proto,
fbe33108 AS	62	struct flow_keys *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
614cd3bd	85	static inline __u64 parse_ipv6(struct __sk_buff skb, __u64 nhoff, __u64 ip_proto,
fbe33108 AS	86	struct flow_keys *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
614cd3bd	99	static inline bool flow_dissector(struct __sk_buff skb, struct flow_keys flow)
fbe33108 AS	100	{
	101	__u64 nhoff = ETH_HLEN;
	102	__u64 ip_proto;
	103	__u64 proto = load_half(skb, 12);
	104	int poff;
	105
	106	if (proto == ETH_P_8021AD) {
	107	proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
	108	h_vlan_encapsulated_proto));
	109	nhoff += sizeof(struct vlan_hdr);
	110	}
	111
	112	if (proto == ETH_P_8021Q) {
	113	proto = load_half(skb, nhoff + offsetof(struct vlan_hdr,
	114	h_vlan_encapsulated_proto));
	115	nhoff += sizeof(struct vlan_hdr);
	116	}
	117
	118	if (likely(proto == ETH_P_IP))
	119	nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
	120	else if (proto == ETH_P_IPV6)
	121	nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
	122	else
	123	return false;
	124
	125	switch (ip_proto) {
	126	case IPPROTO_GRE: {
	127	struct gre_hdr {
	128	__be16 flags;
	129	__be16 proto;
	130	};
	131
	132	__u64 gre_flags = load_half(skb,
	133	nhoff + offsetof(struct gre_hdr, flags));
	134	__u64 gre_proto = load_half(skb,
	135	nhoff + offsetof(struct gre_hdr, proto));
	136
	137	if (gre_flags & (GRE_VERSION\|GRE_ROUTING))
	138	break;
	139
	140	proto = gre_proto;
	141	nhoff += 4;
	142	if (gre_flags & GRE_CSUM)
	143	nhoff += 4;
	144	if (gre_flags & GRE_KEY)
	145	nhoff += 4;
	146	if (gre_flags & GRE_SEQ)
	147	nhoff += 4;
	148
	149	if (proto == ETH_P_8021Q) {
	150	proto = load_half(skb,
	151	nhoff + offsetof(struct vlan_hdr,
	152	h_vlan_encapsulated_proto));
	153	nhoff += sizeof(struct vlan_hdr);
	154	}
	155
	156	if (proto == ETH_P_IP)
	157	nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
	158	else if (proto == ETH_P_IPV6)
	159	nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
	160	else
	161	return false;
	162	break;
	163	}
164	case IPPROTO_IPIP:
165	nhoff = parse_ip(skb, nhoff, &ip_proto, flow);
166	break;
167	case IPPROTO_IPV6:
168	nhoff = parse_ipv6(skb, nhoff, &ip_proto, flow);
169	break;
170	default:
171	break;
172	}
173
174	flow->ip_proto = ip_proto;
175	poff = proto_ports_offset(ip_proto);
176	if (poff >= 0) {
177	nhoff += poff;
178	flow->ports = load_word(skb, nhoff);
179	}
180
181	flow->thoff = (__u16) nhoff;
182
183	return true;
184	}
185
614cd3bd AS	186	struct pair {
	187	long packets;
	188	long bytes;
	189	};
	190
fbe33108 AS	191	struct bpf_map_def SEC("maps") hash_map = {
	192	.type = BPF_MAP_TYPE_HASH,
	193	.key_size = sizeof(__be32),
614cd3bd	194	.value_size = sizeof(struct pair),
fbe33108 AS	195	.max_entries = 1024,
	196	};
	197
	198	SEC("socket2")
614cd3bd	199	int bpf_prog2(struct __sk_buff *skb)
fbe33108 AS	200	{
fbe33108 AS	201	struct flow_keys flow;
614cd3bd	202	struct pair *value;
fbe33108 AS	203	u32 key;
	204
	205	if (!flow_dissector(skb, &flow))
	206	return 0;
	207
	208	key = flow.dst;
	209	value = bpf_map_lookup_elem(&hash_map, &key);
	210	if (value) {
614cd3bd AS	211	__sync_fetch_and_add(&value->packets, 1);
614cd3bd AS	212	__sync_fetch_and_add(&value->bytes, skb->len);
fbe33108	213	} else {
614cd3bd	214	struct pair val = {1, skb->len};
fbe33108 AS	215
	216	bpf_map_update_elem(&hash_map, &key, &val, BPF_ANY);
	217	}
	218	return 0;
	219	}
	220
	221	char _license[] SEC("license") = "GPL";