github.com/cilium/cilium@v1.16.2/bpf/tests/ratelimit.c (about)

     1  // SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause)
     2  /* Copyright Authors of Cilium */
     3  
     4  #include "bpf/ctx/xdp.h"
     5  #include "node_config.h"
     6  #include "common.h"
     7  #include "lib/maps.h"
     8  
     9  static __u64 mock_ktime_get_ns(void)
    10  {
    11  	return 3000 * NSEC_PER_SEC;
    12  }
    13  
    14  #define ktime_get_ns mock_ktime_get_ns
    15  
    16  #include "lib/ratelimit.h"
    17  
    18  CHECK("xdp", "ratelimit") int test_ratelimit(void)
    19  {
    20  	struct ratelimit_settings settings = {
    21  		.bucket_size = 1000,
    22  		.tokens_per_topup = 100,
    23  		.topup_interval_ns = NSEC_PER_SEC,
    24  	};
    25  	struct ratelimit_key key = {
    26  		.netdev_idx = 1,
    27  	};
    28  	struct ratelimit_value *value;
    29  
    30  	test_init();
    31  
    32  	TEST("bucket-created-when-missing", {
    33  		value = map_lookup_elem(&RATELIMIT_MAP, &key);
    34  		if (value)
    35  			test_fatal("Bucket already exits");
    36  
    37  		ratelimit_check_and_take(&key, &settings);
    38  
    39  		value = map_lookup_elem(&RATELIMIT_MAP, &key);
    40  		if (!value)
    41  			test_fatal("Bucket not created");
    42  	})
    43  
    44  	TEST("block-on-bucket-empty", {
    45  		value = map_lookup_elem(&RATELIMIT_MAP, &key);
    46  		if (!value)
    47  			test_fatal("Bucket not created");
    48  
    49  		value->tokens = 1;
    50  		if (!ratelimit_check_and_take(&key, &settings))
    51  			test_fatal("Rate limit not allowed when bucket not empty");
    52  
    53  		if (value->tokens != 0)
    54  			test_fatal("Bucket not empty");
    55  
    56  		if (ratelimit_check_and_take(&key, &settings))
    57  			test_fatal("Rate limit allowed when bucket empty");
    58  	})
    59  
    60  	TEST("topup-after-interval", {
    61  		value = map_lookup_elem(&RATELIMIT_MAP, &key);
    62  		if (!value)
    63  			test_fatal("Bucket not created");
    64  
    65  		/* Set last topup to 1 interval ago */
    66  		value->tokens = 0;
    67  		value->last_topup = ktime_get_ns() - (settings.topup_interval_ns + 1);
    68  
    69  		if (!ratelimit_check_and_take(&key, &settings))
    70  			test_fatal("Rate limit not allowed after topup");
    71  
    72  		if (value->tokens != settings.tokens_per_topup - 1)
    73  			test_fatal("Unexpected token amount after topup");
    74  	})
    75  
    76  	TEST("do-not-go-over-bucket-size", {
    77  		value = map_lookup_elem(&RATELIMIT_MAP, &key);
    78  		if (!value)
    79  			test_fatal("Bucket not created");
    80  
    81  		/* Set last topup to 100 intervals ago */
    82  		value->tokens = 0;
    83  		value->last_topup = ktime_get_ns() - (100 * settings.topup_interval_ns);
    84  
    85  		if (!ratelimit_check_and_take(&key, &settings))
    86  			test_fatal("Rate limit not allowed after topup");
    87  
    88  		if (value->tokens != settings.bucket_size - 1)
    89  			test_fatal("Unexpected token amount after topup");
    90  	})
    91  
    92  	test_finish();
    93  }