bpf/progs/test_tc_edt.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 #include <stdint.h>
0003 #include <linux/bpf.h>
0004 #include <linux/if_ether.h>
0005 #include <linux/stddef.h>
0006 #include <linux/in.h>
0007 #include <linux/ip.h>
0008 #include <linux/pkt_cls.h>
0009 #include <linux/tcp.h>
0010 #include <bpf/bpf_helpers.h>
0011 #include <bpf/bpf_endian.h>
0012
0013 /* the maximum delay we are willing to add (drop packets beyond that) */
0014 #define TIME_HORIZON_NS (2000 * 1000 * 1000)
0015 #define NS_PER_SEC 1000000000
0016 #define ECN_HORIZON_NS 5000000
0017 #define THROTTLE_RATE_BPS (5 * 1000 * 1000)
0018
0019 /* flow_key => last_tstamp timestamp used */
0020 struct {
0021     __uint(type, BPF_MAP_TYPE_HASH);
0022     __type(key, uint32_t);
0023     __type(value, uint64_t);
0024     __uint(max_entries, 1);
0025 } flow_map SEC(".maps");
0026
0027 static inline int throttle_flow(struct __sk_buff *skb)
0028 {
0029     int key = 0;
0030     uint64_t *last_tstamp = bpf_map_lookup_elem(&flow_map, &key);
0031     uint64_t delay_ns = ((uint64_t)skb->len) * NS_PER_SEC /
0032             THROTTLE_RATE_BPS;
0033     uint64_t now = bpf_ktime_get_ns();
0034     uint64_t tstamp, next_tstamp = 0;
0035
0036     if (last_tstamp)
0037         next_tstamp = *last_tstamp + delay_ns;
0038
0039     tstamp = skb->tstamp;
0040     if (tstamp < now)
0041         tstamp = now;
0042
0043     /* should we throttle? */
0044     if (next_tstamp <= tstamp) {
0045         if (bpf_map_update_elem(&flow_map, &key, &tstamp, BPF_ANY))
0046             return TC_ACT_SHOT;
0047         return TC_ACT_OK;
0048     }
0049
0050     /* do not queue past the time horizon */
0051     if (next_tstamp - now >= TIME_HORIZON_NS)
0052         return TC_ACT_SHOT;
0053
0054     /* set ecn bit, if needed */
0055     if (next_tstamp - now >= ECN_HORIZON_NS)
0056         bpf_skb_ecn_set_ce(skb);
0057
0058     if (bpf_map_update_elem(&flow_map, &key, &next_tstamp, BPF_EXIST))
0059         return TC_ACT_SHOT;
0060     skb->tstamp = next_tstamp;
0061
0062     return TC_ACT_OK;
0063 }
0064
0065 static inline int handle_tcp(struct __sk_buff *skb, struct tcphdr *tcp)
0066 {
0067     void *data_end = (void *)(long)skb->data_end;
0068
0069     /* drop malformed packets */
0070     if ((void *)(tcp + 1) > data_end)
0071         return TC_ACT_SHOT;
0072
0073     if (tcp->dest == bpf_htons(9000))
0074         return throttle_flow(skb);
0075
0076     return TC_ACT_OK;
0077 }
0078
0079 static inline int handle_ipv4(struct __sk_buff *skb)
0080 {
0081     void *data_end = (void *)(long)skb->data_end;
0082     void *data = (void *)(long)skb->data;
0083     struct iphdr *iph;
0084     uint32_t ihl;
0085
0086     /* drop malformed packets */
0087     if (data + sizeof(struct ethhdr) > data_end)
0088         return TC_ACT_SHOT;
0089     iph = (struct iphdr *)(data + sizeof(struct ethhdr));
0090     if ((void *)(iph + 1) > data_end)
0091         return TC_ACT_SHOT;
0092     ihl = iph->ihl * 4;
0093     if (((void *)iph) + ihl > data_end)
0094         return TC_ACT_SHOT;
0095
0096     if (iph->protocol == IPPROTO_TCP)
0097         return handle_tcp(skb, (struct tcphdr *)(((void *)iph) + ihl));
0098
0099     return TC_ACT_OK;
0100 }
0101
0102 SEC("cls_test") int tc_prog(struct __sk_buff *skb)
0103 {
0104     if (skb->protocol == bpf_htons(ETH_P_IP))
0105         return handle_ipv4(skb);
0106
0107     return TC_ACT_OK;
0108 }
0109
0110 char __license[] SEC("license") = "GPL";