tools/testing/selftests/bpf/progs/test_tc_edt.c

   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";