From 1b907d13bb97aba8badcb428623fa13e803b8d92 Mon Sep 17 00:00:00 2001
From: trinity-1686a <trinity@deuxfleurs.fr>
Date: Sun, 10 Sep 2023 13:13:11 +0200
Subject: add rate limit on BEGIN and RESOLVE cell per circuit

---
 src/core/or/dos.c | 38 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

(limited to 'src/core/or/dos.c')

diff --git a/src/core/or/dos.c b/src/core/or/dos.c
index a47738c906..63cac190fd 100644
--- a/src/core/or/dos.c
+++ b/src/core/or/dos.c
@@ -325,7 +325,8 @@ get_param_stream_defense_type(const networkstatus_t *ns)
   }
   return networkstatus_get_param(ns, "DoSStreamCreationDefenseType",
                                  DOS_STREAM_DEFENSE_TYPE_DEFAULT,
-                                 DOS_STREAM_DEFENSE_NONE, DOS_STREAM_DEFENSE_MAX);
+                                 DOS_STREAM_DEFENSE_NONE,
+                                 DOS_STREAM_DEFENSE_MAX);
 }
 
 /* Set circuit creation parameters located in the consensus or their default
@@ -836,6 +837,41 @@ dos_conn_addr_get_defense_type(const tor_addr_t *addr)
   return DOS_CONN_DEFENSE_NONE;
 }
 
+/* Stream creation public API. */
+
+/* Return the action to take against a BEGIN or RESOLVE cell. Return
+ *  DOS_STREAM_DEFENSE_NONE when no action should be taken.
+ *  Increment the appropriate counter when the cell was found to go over a
+ *  limit. */
+dos_stream_defense_type_t
+dos_stream_new_begin_or_resolve_cell(or_circuit_t *circ)
+{
+  if (!dos_stream_enabled || circ == NULL)
+    return DOS_STREAM_DEFENSE_NONE;
+
+  token_bucket_ctr_refill(&circ->stream_limiter,
+                          (uint32_t) monotime_coarse_absolute_sec());
+
+  if (token_bucket_ctr_get(&circ->stream_limiter) > 0) {
+    token_bucket_ctr_dec(&circ->stream_limiter, 1);
+    return DOS_STREAM_DEFENSE_NONE;
+  }
+  /* if defense type is DOS_STREAM_DEFENSE_NONE but DoSStreamEnabled is true,
+   * we count offending cells as rejected, despite them being actually
+   * accepted. */
+  ++stream_num_rejected;
+  return dos_stream_defense_type;
+}
+
+/* Initialize the token bucket for stream rate limit on a circuit. */
+void
+dos_stream_init_circ_tbf(or_circuit_t *circ)
+{
+  token_bucket_ctr_init(&circ->stream_limiter, dos_stream_rate,
+                        dos_stream_burst,
+                        (uint32_t) monotime_coarse_absolute_sec());
+}
+
 /* General API */
 
 /* Take any appropriate actions for the given geoip entry that is about to get
-- 
cgit v1.2.3-54-g00ecf