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