3 files changed, 43 insertions, 0 deletions

diff --git a/src/or/dos.c b/src/or/dos.c
index 40e88aead0..5af75ca57d 100644
--- a/src/or/dos.c
+++ b/src/or/dos.c
@@ -528,6 +528,41 @@ dos_conn_addr_get_defense_type(const tor_addr_t *addr)
 
 /* General API */
 
+/* Take any appropriate actions for the given geoip entry that is about to get
+ * freed. This is called for every entry that is being freed.
+ *
+ * This function will clear out the connection tracked flag if the concurrent
+ * count of the entry is above 0 so if those connections end up being seen by
+ * this subsystem, we won't try to decrement the counter for a new geoip entry
+ * that might have been added after this call for the same address. */
+void
+dos_geoip_entry_about_to_free(const clientmap_entry_t *geoip_ent)
+{
+  tor_assert(geoip_ent);
+
+  /* The count is down to 0 meaning no connections right now, we can safely
+   * clear the geoip entry from the cache. */
+  if (geoip_ent->dos_stats.concurrent_count == 0) {
+    goto end;
+  }
+
+  /* For each connection matching the geoip entry address, we'll clear the
+   * tracked flag because the entry is about to get removed from the geoip
+   * cache. We do not try to decrement if the flag is not set. */
+  SMARTLIST_FOREACH_BEGIN(get_connection_array(), connection_t *, conn) {
+    if (conn->type == CONN_TYPE_OR) {
+      or_connection_t *or_conn = TO_OR_CONN(conn);
+      if (!tor_addr_compare(&geoip_ent->addr, &or_conn->real_addr,
+                            CMP_EXACT)) {
+        or_conn->tracked_for_dos_mitigation = 0;
+      }
+    }
+  } SMARTLIST_FOREACH_END(conn);
+
+ end:
+  return;
+}
+
 /* Note down that we've just refused a single hop client. This increments a
  * counter later used for the heartbeat. */
 void
diff --git a/src/or/dos.h b/src/or/dos.h
index 56835169d2..9ce1baddb8 100644
--- a/src/or/dos.h
+++ b/src/or/dos.h
@@ -43,11 +43,15 @@ typedef struct dos_client_stats_t {
 
 /* General API. */
 
+/* Stub. */
+struct clientmap_entry_t;
+
 void dos_init(void);
 void dos_free_all(void);
 void dos_consensus_has_changed(const networkstatus_t *ns);
 int dos_enabled(void);
 void dos_log_heartbeat(void);
+void dos_geoip_entry_about_to_free(const struct clientmap_entry_t *geoip_ent);
 
 void dos_new_client_conn(or_connection_t *or_conn);
 void dos_close_client_conn(const or_connection_t *or_conn);
diff --git a/src/or/geoip.c b/src/or/geoip.c
index 5f0b04b568..4e4f6e639a 100644
--- a/src/or/geoip.c
+++ b/src/or/geoip.c
@@ -516,6 +516,10 @@ clientmap_entry_free(clientmap_entry_t *ent)
   if (!ent)
     return;
 
+  /* This entry is about to be freed so pass it to the DoS subsystem to see if
+   * any actions can be taken about it. */
+  dos_geoip_entry_about_to_free(ent);
+
   tor_free(ent->transport_name);
   tor_free(ent);
 }