1 files changed, 216 insertions, 3 deletions

diff --git a/src/test/test_bwmgt.c b/src/test/test_bwmgt.c
index 5a013aa268..117783cafc 100644
--- a/src/test/test_bwmgt.c
+++ b/src/test/test_bwmgt.c
@@ -1,4 +1,4 @@
-/* Copyright (c) 2018-2019, The Tor Project, Inc. */
+/* Copyright (c) 2018-2020, The Tor Project, Inc. */
 /* See LICENSE for licensing information */
 
 /**
@@ -6,18 +6,70 @@
  * \brief tests for bandwidth management / token bucket functions
  */
 
+#define CONFIG_PRIVATE
+#define CONNECTION_PRIVATE
+#define DIRAUTH_SYS_PRIVATE
 #define TOKEN_BUCKET_PRIVATE
 
 #include "core/or/or.h"
-#include "test/test.h"
 
+#include "app/config/config.h"
+#include "core/mainloop/connection.h"
+#include "feature/dirauth/dirauth_sys.h"
+#include "feature/dircommon/directory.h"
+#include "feature/nodelist/microdesc.h"
+#include "feature/nodelist/networkstatus.h"
+#include "feature/nodelist/nodelist.h"
+#include "feature/nodelist/routerlist.h"
+#include "lib/crypt_ops/crypto_rand.h"
 #include "lib/evloop/token_bucket.h"
+#include "test/test.h"
+#include "test/test_helpers.h"
+
+#include "app/config/or_options_st.h"
+#include "core/or/connection_st.h"
+#include "feature/dirauth/dirauth_options_st.h"
+#include "feature/nodelist/microdesc_st.h"
+#include "feature/nodelist/networkstatus_st.h"
+#include "feature/nodelist/routerinfo_st.h"
+#include "feature/nodelist/routerstatus_st.h"
 
 // an imaginary time, in timestamp units. Chosen so it will roll over.
 static const uint32_t START_TS = UINT32_MAX-10;
 static const int32_t KB = 1024;
 static const uint32_t GB = (UINT64_C(1) << 30);
 
+static or_options_t mock_options;
+
+static const or_options_t *
+mock_get_options(void)
+{
+  return &mock_options;
+}
+
+static networkstatus_t *dummy_ns = NULL;
+static networkstatus_t *
+mock_networkstatus_get_latest_consensus(void)
+{
+  return dummy_ns;
+}
+
+static networkstatus_t *
+mock_networkstatus_get_latest_consensus_by_flavor(consensus_flavor_t f)
+{
+  tor_assert(f == FLAV_MICRODESC);
+  return dummy_ns;
+}
+
+/* Number of address a single node_t can have. Default to the production
+ * value. This is to control the size of the bloom filter. */
+static int addr_per_node = 2;
+static int
+mock_get_estimated_address_per_node(void)
+{
+  return addr_per_node;
+}
+
 static void
 test_bwmgt_token_buf_init(void *arg)
 {
@@ -220,8 +272,167 @@ test_bwmgt_token_buf_helpers(void *arg)
   ;
 }
 
+static void
+test_bwmgt_dir_conn_global_write_low(void *arg)
+{
+  bool ret;
+  int addr_family;
+  connection_t *conn = NULL;
+  routerstatus_t *rs = NULL; microdesc_t *md = NULL; routerinfo_t *ri = NULL;
+  tor_addr_t relay_addr;
+  dirauth_options_t *dirauth_opts = NULL;
+
+  (void) arg;
+
+  memset(&mock_options, 0, sizeof(or_options_t));
+  MOCK(networkstatus_get_latest_consensus,
+       mock_networkstatus_get_latest_consensus);
+  MOCK(networkstatus_get_latest_consensus_by_flavor,
+       mock_networkstatus_get_latest_consensus_by_flavor);
+  MOCK(get_estimated_address_per_node,
+       mock_get_estimated_address_per_node);
+
+  /*
+   * The following is rather complex but that is what it takes to add a dummy
+   * consensus with a valid routerlist which will populate our node address
+   * set that we need to lookup to test the known relay code path.
+   *
+   * We MUST do that before we MOCK(get_options) else it is another world of
+   * complexity.
+   */
+
+  /* This will be the address of our relay. */
+  tor_addr_parse(&relay_addr, "1.2.3.4");
+
+  /* We'll now add a relay into our routerlist and see if we let it. */
+  dummy_ns = tor_malloc_zero(sizeof(*dummy_ns));
+  dummy_ns->flavor = FLAV_MICRODESC;
+  dummy_ns->routerstatus_list = smartlist_new();
+
+  md = tor_malloc_zero(sizeof(*md));
+  ri = tor_malloc_zero(sizeof(*ri));
+  rs = tor_malloc_zero(sizeof(*rs));
+  crypto_rand(rs->identity_digest, sizeof(rs->identity_digest));
+  crypto_rand(md->digest, sizeof(md->digest));
+  memcpy(rs->descriptor_digest, md->digest, DIGEST256_LEN);
+
+  /* Set IP address. */
+  rs->addr = tor_addr_to_ipv4h(&relay_addr);
+  ri->addr = rs->addr;
+  /* Add the rs to the consensus becoming a node_t. */
+  smartlist_add(dummy_ns->routerstatus_list, rs);
+
+  /* Add all configured authorities (hardcoded) before we set the consensus so
+   * the address set exists. */
+  ret = consider_adding_dir_servers(&mock_options, &mock_options);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* This will make the nodelist bloom filter very large
+   * (the_nodelist->node_addrs) so we will fail the contain test rarely. */
+  addr_per_node = 1024;
+
+  nodelist_set_consensus(dummy_ns);
+
+  dirauth_opts = tor_malloc_zero(sizeof(dirauth_options_t));
+  dirauth_opts->AuthDirRejectRequestsUnderLoad = 0;
+  dirauth_set_options(dirauth_opts);
+
+  /* Ok, now time to control which options we use. */
+  MOCK(get_options, mock_get_options);
+
+  /* Set ourselves as an authoritative dir. */
+  mock_options.AuthoritativeDir = 1;
+  mock_options.V3AuthoritativeDir = 1;
+  mock_options.UseDefaultFallbackDirs = 0;
+
+  /* This will set our global bucket to 1 byte and thus we will hit the
+   * banwdith limit in our test. */
+  mock_options.BandwidthRate = 1;
+  mock_options.BandwidthBurst = 1;
+
+  /* Else an IPv4 address screams. */
+  mock_options.ClientUseIPv4 = 1;
+  mock_options.ClientUseIPv6 = 1;
+
+  /* Initialize the global buckets. */
+  connection_bucket_init();
+
+  /* The address "127.0.0.1" is set with this helper. */
+  conn = test_conn_get_connection(DIR_CONN_STATE_MIN_, CONN_TYPE_DIR,
+                                  DIR_PURPOSE_MIN_);
+  tt_assert(conn);
+
+  /* First try a non authority non relay IP thus a client but we are not
+   * configured to reject requests under load so we should get a false value
+   * that our limit is _not_ low. */
+  addr_family = tor_addr_parse(&conn->addr, "1.1.1.1");
+  tt_int_op(addr_family, OP_EQ, AF_INET);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* Now, we will reject requests under load so try again a non authority non
+   * relay IP thus a client. We should get a warning that our limit is too
+   * low. */
+  dirauth_opts->AuthDirRejectRequestsUnderLoad = 1;
+
+  addr_family = tor_addr_parse(&conn->addr, "1.1.1.1");
+  tt_int_op(addr_family, OP_EQ, AF_INET);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 1);
+
+  /* Now, lets try with a connection address from moria1. It should always
+   * pass even though our limit is too low. */
+  addr_family = tor_addr_parse(&conn->addr, "128.31.0.39");
+  tt_int_op(addr_family, OP_EQ, AF_INET);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* IPv6 testing of gabelmoo. */
+  addr_family = tor_addr_parse(&conn->addr, "[2001:638:a000:4140::ffff:189]");
+  tt_int_op(addr_family, OP_EQ, AF_INET6);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* Lets retry with a known relay address. It should pass. Possible due to
+   * our consensus setting above. */
+  memcpy(&conn->addr, &relay_addr, sizeof(tor_addr_t));
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* Lets retry with a random IP that is not an authority nor a relay. */
+  addr_family = tor_addr_parse(&conn->addr, "1.2.3.4");
+  tt_int_op(addr_family, OP_EQ, AF_INET);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 0);
+
+  /* Finally, just make sure it still denies an IP if we are _not_ a v3
+   * directory authority. */
+  mock_options.V3AuthoritativeDir = 0;
+  addr_family = tor_addr_parse(&conn->addr, "1.2.3.4");
+  tt_int_op(addr_family, OP_EQ, AF_INET);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 1);
+
+  /* Random IPv6 should not be allowed. */
+  addr_family = tor_addr_parse(&conn->addr, "[CAFE::ACAB]");
+  tt_int_op(addr_family, OP_EQ, AF_INET6);
+  ret = connection_dir_is_global_write_low(conn, INT_MAX);
+  tt_int_op(ret, OP_EQ, 1);
+
+ done:
+  connection_free_minimal(conn);
+  routerstatus_free(rs); routerinfo_free(ri); microdesc_free(md);
+  smartlist_clear(dummy_ns->routerstatus_list);
+  networkstatus_vote_free(dummy_ns);
+
+  UNMOCK(get_estimated_address_per_node);
+  UNMOCK(networkstatus_get_latest_consensus);
+  UNMOCK(networkstatus_get_latest_consensus_by_flavor);
+  UNMOCK(get_options);
+}
+
 #define BWMGT(name)                                          \
-  { #name, test_bwmgt_ ## name , 0, NULL, NULL }
+  { #name, test_bwmgt_ ## name , TT_FORK, NULL, NULL }
 
 struct testcase_t bwmgt_tests[] = {
   BWMGT(token_buf_init),
@@ -229,5 +440,7 @@ struct testcase_t bwmgt_tests[] = {
   BWMGT(token_buf_dec),
   BWMGT(token_buf_refill),
   BWMGT(token_buf_helpers),
+
+  BWMGT(dir_conn_global_write_low),
   END_OF_TESTCASES
 };