tjh.dev
Linux kernel mirror (for testing)
git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git
kernel
os
linux
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 */
19struct 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
26static 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
64static 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
78static 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
101SEC("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
109char __license[] SEC("license") = "GPL";