From aac898343ee52e77f97fe73f0cda4d73ec926814 Mon Sep 17 00:00:00 2001
From: Miroslav Lichvar <mlichvar@redhat.com>
Date: Tue, 2 Apr 2024 15:05:19 +0200
Subject: [PATCH] clientlog: add support for KoD rate limiting

Add a third return value to CLG_LimitServiceRate() to indicate the
server should send a response requesting the client to reduce its
polling rate. It randomly selects from a fraction (configurable to 1/2,
1/4, 1/8, 1/16, or disabled) of responses which would be dropped
(after selecting responses for the leak option).
---
 clientlog.c           | 28 +++++++++++++++------
 clientlog.h           |  1 +
 test/unit/clientlog.c | 58 +++++++++++++++++++++++++++++++++----------
 3 files changed, 67 insertions(+), 20 deletions(-)

diff --git a/clientlog.c b/clientlog.c
index 5c8a981..43cba67 100644
--- a/clientlog.c
+++ b/clientlog.c
@@ -117,6 +117,14 @@ static int token_shift[MAX_SERVICES];
 
 static int leak_rate[MAX_SERVICES];
 
+/* Rates at which responses requesting clients to reduce their rate
+   (e.g. NTP KoD RATE) are randomly allowed (in log2, but 0 means disabled) */
+
+#define MIN_KOD_RATE 0
+#define MAX_KOD_RATE 4
+
+static int kod_rate[MAX_SERVICES];
+
 /* Limit intervals in log2 */
 static int limit_interval[MAX_SERVICES];
 
@@ -354,13 +362,14 @@ set_bucket_params(int interval, int burst, uint16_t *max_tokens,
 void
 CLG_Initialise(void)
 {
-  int i, interval, burst, lrate, slots2;
+  int i, interval, burst, lrate, krate, slots2;
 
   for (i = 0; i < MAX_SERVICES; i++) {
     max_tokens[i] = 0;
     tokens_per_hit[i] = 0;
     token_shift[i] = 0;
     leak_rate[i] = 0;
+    kod_rate[i] = 0;
     limit_interval[i] = MIN_LIMIT_INTERVAL;
 
     switch (i) {
@@ -382,6 +391,7 @@ CLG_Initialise(void)
 
     set_bucket_params(interval, burst, &max_tokens[i], &tokens_per_hit[i], &token_shift[i]);
     leak_rate[i] = CLAMP(MIN_LEAK_RATE, lrate, MAX_LEAK_RATE);
+    kod_rate[i] = CLAMP(MIN_KOD_RATE, krate, MAX_KOD_RATE);
     limit_interval[i] = CLAMP(MIN_LIMIT_INTERVAL, interval, MAX_LIMIT_INTERVAL);
   }
 
@@ -579,21 +589,21 @@ CLG_LogServiceAccess(CLG_Service service, IPAddr *client, struct timespec *now)
 /* ================================================== */
 
 static int
-limit_response_random(int leak_rate)
+limit_response_random(int rate)
 {
   static uint32_t rnd;
   static int bits_left = 0;
   int r;
 
-  if (bits_left < leak_rate) {
+  if (bits_left < rate) {
     UTI_GetRandomBytes(&rnd, sizeof (rnd));
     bits_left = 8 * sizeof (rnd);
   }
 
-  /* Return zero on average once per 2^leak_rate */
-  r = rnd % (1U << leak_rate) ? 1 : 0;
-  rnd >>= leak_rate;
-  bits_left -= leak_rate;
+  /* Return zero on average once per 2^rate */
+  r = rnd % (1U << rate) ? 1 : 0;
+  rnd >>= rate;
+  bits_left -= rate;
 
   return r;
 }
@@ -635,6 +645,10 @@ CLG_LimitServiceRate(CLG_Service service, int index)
     return CLG_PASS;
   }
 
+  if (kod_rate[service] > 0 && !limit_response_random(kod_rate[service])) {
+    return CLG_KOD;
+  }
+
   record->drop_flags |= 1U << service;
   record->drops[service]++;
   total_drops[service]++;
diff --git a/clientlog.h b/clientlog.h
index 0d7df00..2e4d99c 100644
--- a/clientlog.h
+++ b/clientlog.h
@@ -40,6 +40,7 @@ typedef enum {
 typedef enum {
   CLG_PASS = 0,
   CLG_DROP,
+  CLG_KOD,
 } CLG_Limit;
 
 extern void CLG_Initialise(void);
diff --git a/test/unit/clientlog.c b/test/unit/clientlog.c
index 59ec2b7..f08b949 100644
--- a/test/unit/clientlog.c
+++ b/test/unit/clientlog.c
@@ -35,18 +35,18 @@ void
 test_unit(void)
 {
   uint64_t ts64, prev_first_ts64, prev_last_ts64, max_step;
+  int i, j, k, kod, passes, kods, drops, index, shift;
   uint32_t index2, prev_first, prev_size;
   NTP_Timestamp_Source ts_src, ts_src2;
   struct timespec ts, ts2;
-  int i, j, k, index, shift;
   CLG_Service s;
   NTP_int64 ntp_ts;
   IPAddr ip;
   char conf[][100] = {
     "clientloglimit 20000",
     "ratelimit interval 3 burst 4 leak 3",
-    "cmdratelimit interval 3 burst 4 leak 3",
-    "ntsratelimit interval 6 burst 8 leak 3",
+    "ntsratelimit interval 4 burst 8 leak 3",
+    "cmdratelimit interval 6 burst 4 leak 3",
   };
 
   CNF_Initialise(0, 0);
@@ -80,19 +80,51 @@ test_unit(void)
   DEBUG_LOG("records %u", ARR_GetSize(records));
   TEST_CHECK(ARR_GetSize(records) == 128);
 
-  s = CLG_NTP;
+  for (kod = 0; kod <= 2; kod += 2) {
+    for (s = CLG_NTP; s <= CLG_CMDMON; s++) {
+      for (i = passes = kods = drops = 0; i < 10000; i++) {
+        kod_rate[s] = kod;
+        ts.tv_sec += 1;
+        index = CLG_LogServiceAccess(s, &ip, &ts);
+        TEST_CHECK(index >= 0);
+        switch (CLG_LimitServiceRate(s, index)) {
+          case CLG_PASS:
+            passes += 1;
+            break;
+          case CLG_DROP:
+            drops += 1;
+            break;
+          case CLG_KOD:
+            kods += 1;
+            break;
+          default:
+            assert(0);
+        }
+      }
 
-  for (i = j = 0; i < 10000; i++) {
-    ts.tv_sec += 1;
-    index = CLG_LogServiceAccess(s, &ip, &ts);
-    TEST_CHECK(index >= 0);
-    if (CLG_LimitServiceRate(s, index) == CLG_PASS)
-      j++;
+      DEBUG_LOG("service %d requests %d passes %d kods %d drops %d",
+                (int)s, i, passes, kods, drops);
+      if (kod)
+        TEST_CHECK(kods * 2.5 < drops && kods * 3.5 > drops);
+      else
+        TEST_CHECK(kods == 0);
+
+      switch (s) {
+        case CLG_NTP:
+          TEST_CHECK(passes > 1750 && passes < 2050);
+          break;
+        case CLG_NTSKE:
+          TEST_CHECK(passes > 1300 && passes < 1600);
+          break;
+        case CLG_CMDMON:
+          TEST_CHECK(passes > 1100 && passes < 1400);
+          break;
+        default:
+          assert(0);
+      }
+    }
   }
 
-  DEBUG_LOG("requests %d responses %d", i, j);
-  TEST_CHECK(j * 4 < i && j * 6 > i);
-
   TEST_CHECK(!ntp_ts_map.timestamps);
 
   UTI_ZeroNtp64(&ntp_ts);