1 // SPDX-License-Identifier: GPL-2.0
2 #include <stdint.h>
3 #include <linux/bpf.h>
4 #include <linux/if_ether.h>
5 #include <linux/in.h>
6 #include <linux/ip.h>
7 #include <linux/pkt_cls.h>
8 #include <linux/tcp.h>
9 #include "bpf_helpers.h"
10 #include "bpf_endian.h"
11 
12 /* the maximum delay we are willing to add (drop packets beyond that) */
13 #define TIME_HORIZON_NS (2000 * 1000 * 1000)
14 #define NS_PER_SEC 1000000000
15 #define ECN_HORIZON_NS 5000000
16 #define THROTTLE_RATE_BPS (5 * 1000 * 1000)
17 
18 /* flow_key => last_tstamp timestamp used */
19 struct bpf_map_def SEC("maps") flow_map = {
20 	.type = BPF_MAP_TYPE_HASH,
21 	.key_size = sizeof(uint32_t),
22 	.value_size = sizeof(uint64_t),
23 	.max_entries = 1,
24 };
25 
26 static inline int throttle_flow(struct __sk_buff *skb)
27 {
28 	int key = 0;
29 	uint64_t *last_tstamp = bpf_map_lookup_elem(&flow_map, &key);
30 	uint64_t delay_ns = ((uint64_t)skb->len) * NS_PER_SEC /
31 			THROTTLE_RATE_BPS;
32 	uint64_t now = bpf_ktime_get_ns();
33 	uint64_t tstamp, next_tstamp = 0;
34 
35 	if (last_tstamp)
36 		next_tstamp = *last_tstamp + delay_ns;
37 
38 	tstamp = skb->tstamp;
39 	if (tstamp < now)
40 		tstamp = now;
41 
42 	/* should we throttle? */
43 	if (next_tstamp <= tstamp) {
44 		if (bpf_map_update_elem(&flow_map, &key, &tstamp, BPF_ANY))
45 			return TC_ACT_SHOT;
46 		return TC_ACT_OK;
47 	}
48 
49 	/* do not queue past the time horizon */
50 	if (next_tstamp - now >= TIME_HORIZON_NS)
51 		return TC_ACT_SHOT;
52 
53 	/* set ecn bit, if needed */
54 	if (next_tstamp - now >= ECN_HORIZON_NS)
55 		bpf_skb_ecn_set_ce(skb);
56 
57 	if (bpf_map_update_elem(&flow_map, &key, &next_tstamp, BPF_EXIST))
58 		return TC_ACT_SHOT;
59 	skb->tstamp = next_tstamp;
60 
61 	return TC_ACT_OK;
62 }
63 
64 static inline int handle_tcp(struct __sk_buff *skb, struct tcphdr *tcp)
65 {
66 	void *data_end = (void *)(long)skb->data_end;
67 
68 	/* drop malformed packets */
69 	if ((void *)(tcp + 1) > data_end)
70 		return TC_ACT_SHOT;
71 
72 	if (tcp->dest == bpf_htons(9000))
73 		return throttle_flow(skb);
74 
75 	return TC_ACT_OK;
76 }
77 
78 static inline int handle_ipv4(struct __sk_buff *skb)
79 {
80 	void *data_end = (void *)(long)skb->data_end;
81 	void *data = (void *)(long)skb->data;
82 	struct iphdr *iph;
83 	uint32_t ihl;
84 
85 	/* drop malformed packets */
86 	if (data + sizeof(struct ethhdr) > data_end)
87 		return TC_ACT_SHOT;
88 	iph = (struct iphdr *)(data + sizeof(struct ethhdr));
89 	if ((void *)(iph + 1) > data_end)
90 		return TC_ACT_SHOT;
91 	ihl = iph->ihl * 4;
92 	if (((void *)iph) + ihl > data_end)
93 		return TC_ACT_SHOT;
94 
95 	if (iph->protocol == IPPROTO_TCP)
96 		return handle_tcp(skb, (struct tcphdr *)(((void *)iph) + ihl));
97 
98 	return TC_ACT_OK;
99 }
100 
101 SEC("cls_test") int tc_prog(struct __sk_buff *skb)
102 {
103 	if (skb->protocol == bpf_htons(ETH_P_IP))
104 		return handle_ipv4(skb);
105 
106 	return TC_ACT_OK;
107 }
108 
109 char __license[] SEC("license") = "GPL";
110