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 }