wg_cookie: make ratelimiter global

Signed-off-by: Matt Dunwoodie <ncon@noconroy.net>

5 files changed, 113 insertions, 115 deletions

diff --git a/TODO.md b/TODO.md
index 9b5ffc3..5f68380 100644
--- a/TODO.md
+++ b/TODO.md
@@ -17,10 +17,9 @@
 - Handle failures of `rn_inithead` and remember to call `rn_detachhead`
   somewhere during cleanup.
 - Stop using `M_WAITOK` and use `M_NOWAIT` instead.
-- Make sure ratelimiter is empty and deinited.
 - Check return value of `rn_inithead`.
 - Perhaps call `rn_detachhead` to free memory when destroying aip.
-- Have one rate limiter table per module, and hash in jail/fib pointer.
+- Hash in jail/fib pointer to ratelimiter.
 
 ### Crypto TODO
 
diff --git a/src/if_wg.c b/src/if_wg.c
index e6fcfaa..e898caf 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -287,7 +287,6 @@ struct wg_softc {
 
 static int clone_count;
 static uma_zone_t wg_packet_zone;
-static uma_zone_t ratelimit_zone;
 static volatile unsigned long peer_counter = 0;
 static const char wgname[] = "wg";
 static unsigned wg_osd_jail_slot;
@@ -2617,8 +2616,7 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 		return (ENOMEM);
 	}
 
-	/* TODO check checker_init return value */
-	cookie_checker_init(&sc->sc_cookie, ratelimit_zone);
+	cookie_checker_init(&sc->sc_cookie);
 
 	sc->sc_socket.so_port = 0;
 
@@ -2834,10 +2832,10 @@ wg_module_init(void)
 
 	wg_packet_zone = uma_zcreate("wg packet", sizeof(struct wg_packet),
 	     NULL, NULL, NULL, NULL, 0, 0);
-	ratelimit_zone = uma_zcreate("wg ratelimit", sizeof(struct ratelimit),
-	     NULL, NULL, NULL, NULL, 0, 0);
 	wg_osd_jail_slot = osd_jail_register(NULL, methods);
 
+	cookie_init();
+
 	wg_run_selftests();
 }
 
@@ -2846,8 +2844,8 @@ wg_module_deinit(void)
 {
 
 	uma_zdestroy(wg_packet_zone);
-	uma_zdestroy(ratelimit_zone);
 	osd_jail_deregister(wg_osd_jail_slot);
+	cookie_deinit();
 
 	MPASS(LIST_EMPTY(&wg_list));
 }
diff --git a/src/selftest/cookie.c b/src/selftest/cookie.c
index c24022e..d6bca19 100644
--- a/src/selftest/cookie.c
+++ b/src/selftest/cookie.c
@@ -29,7 +29,7 @@ static const struct expected_results {
 };
 
 static void
-cookie_ratelimit_timings_test(uma_zone_t zone)
+cookie_ratelimit_timings_test(void)
 {
 	struct ratelimit rl;
 	struct sockaddr_in sin;
@@ -38,7 +38,7 @@ cookie_ratelimit_timings_test(uma_zone_t zone)
 #endif
 	int i;
 
-	ratelimit_init(&rl, zone);
+	ratelimit_init(&rl);
 
 	sin.sin_family = AF_INET;
 #ifdef INET6
@@ -94,13 +94,13 @@ cleanup:
 }
 
 static void
-cookie_ratelimit_capacity_test(uma_zone_t zone)
+cookie_ratelimit_capacity_test(void)
 {
 	struct ratelimit rl;
 	struct sockaddr_in sin;
 	int i;
 
-	ratelimit_init(&rl, zone);
+	ratelimit_init(&rl);
 
 	sin.sin_family = AF_INET;
 	sin.sin_port = 1234;
@@ -123,13 +123,13 @@ cleanup:
 }
 
 static void
-cookie_ratelimit_gc_test(uma_zone_t zone)
+cookie_ratelimit_gc_test(void)
 {
 	struct ratelimit rl;
 	struct sockaddr_in sin;
 	int i;
 
-	ratelimit_init(&rl, zone);
+	ratelimit_init(&rl);
 
 	sin.sin_family = AF_INET;
 	sin.sin_port = 1234;
@@ -168,7 +168,7 @@ cleanup:
 }
 
 static void
-cookie_mac_test(uma_zone_t zone)
+cookie_mac_test(void)
 {
 	struct cookie_checker checker;
 	struct cookie_maker maker;
@@ -187,8 +187,7 @@ cookie_mac_test(uma_zone_t zone)
 	/* Init cookie_maker. */
 	cookie_maker_init(&maker, shared);
 
-	if (cookie_checker_init(&checker, zone) != 0)
-		T_FAILED("cookie_checker_allocate");
+	cookie_checker_init(&checker);
 	cookie_checker_update(&checker, shared);
 
 	/* Create dummy sockaddr */
@@ -280,21 +279,14 @@ cookie_mac_test(uma_zone_t zone)
 
 	T_PASSED;
 cleanup:
-	cookie_checker_deinit(&checker);
+	return;
 }
 
 void
 cookie_selftest(void)
 {
-	uma_zone_t rl_zone;
-
-	rl_zone = uma_zcreate("cookie test", sizeof(struct ratelimit),
-	    NULL, NULL, NULL, NULL, 0, 0);
-
-	cookie_ratelimit_timings_test(rl_zone);
-	cookie_ratelimit_capacity_test(rl_zone);
-	cookie_ratelimit_gc_test(rl_zone);
-	cookie_mac_test(rl_zone);
-
-	uma_zdestroy(rl_zone);
+	cookie_ratelimit_timings_test();
+	cookie_ratelimit_capacity_test();
+	cookie_ratelimit_gc_test();
+	cookie_mac_test();
 }
diff --git a/src/wg_cookie.c b/src/wg_cookie.c
index 94a8aae..c9d3676 100644
--- a/src/wg_cookie.c
+++ b/src/wg_cookie.c
@@ -15,6 +15,45 @@
 #include "support.h"
 #include "wg_cookie.h"
 
+#define COOKIE_MAC1_KEY_LABEL	"mac1----"
+#define COOKIE_COOKIE_KEY_LABEL	"cookie--"
+#define COOKIE_SECRET_MAX_AGE	120
+#define COOKIE_SECRET_LATENCY	5
+
+/* Constants for initiation rate limiting */
+#define RATELIMIT_SIZE		(1 << 13)
+#define RATELIMIT_SIZE_MAX	(RATELIMIT_SIZE * 8)
+#define INITIATIONS_PER_SECOND	20
+#define INITIATIONS_BURSTABLE	5
+#define INITIATION_COST		(SBT_1S / INITIATIONS_PER_SECOND)
+#define TOKEN_MAX		(INITIATION_COST * INITIATIONS_BURSTABLE)
+#define ELEMENT_TIMEOUT		1
+#define IPV4_MASK_SIZE		4 /* Use all 4 bytes of IPv4 address */
+#define IPV6_MASK_SIZE		8 /* Use top 8 bytes (/64) of IPv6 address */
+
+struct ratelimit_entry {
+	LIST_ENTRY(ratelimit_entry)	 r_entry;
+	sa_family_t			 r_af;
+	union {
+		struct in_addr		 r_in;
+#ifdef INET6
+		struct in6_addr		 r_in6;
+#endif
+	};
+	sbintime_t			 r_last_time;	/* sbinuptime */
+	uint64_t			 r_tokens;
+};
+
+struct ratelimit {
+	uint8_t				 rl_secret[SIPHASH_KEY_LENGTH];
+
+	struct rwlock			 rl_lock;
+	struct callout			 rl_gc;
+	LIST_HEAD(, ratelimit_entry)	*rl_table;
+	u_long				 rl_table_mask;
+	size_t				 rl_table_num;
+};
+
 static void	cookie_precompute_key(uint8_t *,
 			const uint8_t[COOKIE_INPUT_SIZE], const char *);
 static void	cookie_macs_mac1(struct cookie_macs *, const void *, size_t,
@@ -24,7 +63,7 @@ static void	cookie_macs_mac2(struct cookie_macs *, const void *, size_t,
 static int	cookie_timer_expired(sbintime_t, uint32_t, uint32_t);
 static void	cookie_checker_make_cookie(struct cookie_checker *,
 			uint8_t[COOKIE_COOKIE_SIZE], struct sockaddr *);
-static int	ratelimit_init(struct ratelimit *, uma_zone_t);
+static int	ratelimit_init(struct ratelimit *);
 static void	ratelimit_deinit(struct ratelimit *);
 static void	ratelimit_gc_callout(void *);
 static void	ratelimit_gc_schedule(struct ratelimit *);
@@ -32,30 +71,26 @@ static void	ratelimit_gc(struct ratelimit *, bool);
 static int	ratelimit_allow(struct ratelimit *, struct sockaddr *);
 static uint64_t siphash13(const uint8_t [SIPHASH_KEY_LENGTH], const void *, size_t);
 
-/* Public Functions */
-void
-cookie_maker_init(struct cookie_maker *cp, const uint8_t key[COOKIE_INPUT_SIZE])
-{
-	bzero(cp, sizeof(*cp));
-	cookie_precompute_key(cp->cp_mac1_key, key, COOKIE_MAC1_KEY_LABEL);
-	cookie_precompute_key(cp->cp_cookie_key, key, COOKIE_COOKIE_KEY_LABEL);
-	rw_init(&cp->cp_lock, "cookie_maker");
-}
+static struct ratelimit ratelimit_v4;
+#ifdef INET6
+static struct ratelimit ratelimit_v6;
+#endif
+static uma_zone_t ratelimit_zone;
 
+/* Public Functions */
 int
-cookie_checker_init(struct cookie_checker *cc, uma_zone_t zone)
+cookie_init(void)
 {
 	int res;
-	bzero(cc, sizeof(*cc));
 
-	rw_init(&cc->cc_key_lock, "cookie_checker_key");
-	rw_init(&cc->cc_secret_lock, "cookie_checker_secret");
+	ratelimit_zone = uma_zcreate("wg ratelimit", sizeof(struct ratelimit),
+	     NULL, NULL, NULL, NULL, 0, 0);
 
-	if ((res = ratelimit_init(&cc->cc_ratelimit_v4, zone)) != 0)
+	if ((res = ratelimit_init(&ratelimit_v4)) != 0)
 		return res;
 #ifdef INET6
-	if ((res = ratelimit_init(&cc->cc_ratelimit_v6, zone)) != 0) {
-		ratelimit_deinit(&cc->cc_ratelimit_v4);
+	if ((res = ratelimit_init(&ratelimit_v6)) != 0) {
+		ratelimit_deinit(&ratelimit_v4);
 		return res;
 	}
 #endif
@@ -63,6 +98,25 @@ cookie_checker_init(struct cookie_checker *cc, uma_zone_t zone)
 }
 
 void
+cookie_deinit(void)
+{
+	uma_zdestroy(ratelimit_zone);
+	ratelimit_deinit(&ratelimit_v4);
+#ifdef INET6
+	ratelimit_deinit(&ratelimit_v6);
+#endif
+}
+
+void
+cookie_checker_init(struct cookie_checker *cc)
+{
+	bzero(cc, sizeof(*cc));
+
+	rw_init(&cc->cc_key_lock, "cookie_checker_key");
+	rw_init(&cc->cc_secret_lock, "cookie_checker_secret");
+}
+
+void
 cookie_checker_update(struct cookie_checker *cc,
     const uint8_t key[COOKIE_INPUT_SIZE])
 {
@@ -78,15 +132,6 @@ cookie_checker_update(struct cookie_checker *cc,
 }
 
 void
-cookie_checker_deinit(struct cookie_checker *cc)
-{
-	ratelimit_deinit(&cc->cc_ratelimit_v4);
-#ifdef INET6
-	ratelimit_deinit(&cc->cc_ratelimit_v6);
-#endif
-}
-
-void
 cookie_checker_create_payload(struct cookie_checker *cc,
     struct cookie_macs *cm, uint8_t nonce[COOKIE_NONCE_SIZE],
     uint8_t ecookie[COOKIE_ENCRYPTED_SIZE], struct sockaddr *sa)
@@ -104,6 +149,15 @@ cookie_checker_create_payload(struct cookie_checker *cc,
 	explicit_bzero(cookie, sizeof(cookie));
 }
 
+void
+cookie_maker_init(struct cookie_maker *cp, const uint8_t key[COOKIE_INPUT_SIZE])
+{
+	bzero(cp, sizeof(*cp));
+	cookie_precompute_key(cp->cp_mac1_key, key, COOKIE_MAC1_KEY_LABEL);
+	cookie_precompute_key(cp->cp_cookie_key, key, COOKIE_COOKIE_KEY_LABEL);
+	rw_init(&cp->cp_lock, "cookie_maker");
+}
+
 int
 cookie_maker_consume_payload(struct cookie_maker *cp,
     uint8_t nonce[COOKIE_NONCE_SIZE], uint8_t ecookie[COOKIE_ENCRYPTED_SIZE])
@@ -182,10 +236,10 @@ cookie_checker_validate_macs(struct cookie_checker *cc, struct cookie_macs *cm,
 		 * implying there is no ratelimiting, or we should ratelimit
 		 * (refuse) respectively. */
 		if (sa->sa_family == AF_INET)
-			return ratelimit_allow(&cc->cc_ratelimit_v4, sa);
+			return ratelimit_allow(&ratelimit_v4, sa);
 #ifdef INET6
 		else if (sa->sa_family == AF_INET6)
-			return ratelimit_allow(&cc->cc_ratelimit_v6, sa);
+			return ratelimit_allow(&ratelimit_v6, sa);
 #endif
 		else
 			return EAFNOSUPPORT;
@@ -272,14 +326,13 @@ cookie_checker_make_cookie(struct cookie_checker *cc,
 }
 
 static int
-ratelimit_init(struct ratelimit *rl, uma_zone_t zone)
+ratelimit_init(struct ratelimit *rl)
 {
 	rw_init(&rl->rl_lock, "ratelimit_lock");
 	callout_init_rw(&rl->rl_gc, &rl->rl_lock, 0);
 	arc4random_buf(rl->rl_secret, sizeof(rl->rl_secret));
 	rl->rl_table = hashinit_flags(RATELIMIT_SIZE, M_DEVBUF,
 	    &rl->rl_table_mask, M_NOWAIT);
-	rl->rl_zone = zone;
 	rl->rl_table_num = 0;
 	return rl->rl_table == NULL ? ENOBUFS : 0;
 }
@@ -334,7 +387,7 @@ ratelimit_gc(struct ratelimit *rl, bool force)
 			if (r->r_last_time < expiry || force) {
 				rl->rl_table_num--;
 				LIST_REMOVE(r, r_entry);
-				uma_zfree(rl->rl_zone, r);
+				uma_zfree(ratelimit_zone, r);
 			}
 		}
 	}
@@ -408,7 +461,7 @@ ratelimit_allow(struct ratelimit *rl, struct sockaddr *sa)
 		goto error;
 
 	/* Goto error if out of memory */
-	if ((r = uma_zalloc(rl->rl_zone, M_NOWAIT)) == NULL)
+	if ((r = uma_zalloc(ratelimit_zone, M_NOWAIT)) == NULL)
 		goto error;
 
 	rl->rl_table_num++;
diff --git a/src/wg_cookie.h b/src/wg_cookie.h
index b86831b..3ffa7aa 100644
--- a/src/wg_cookie.h
+++ b/src/wg_cookie.h
@@ -23,51 +23,11 @@
 #define COOKIE_INPUT_SIZE	32
 #define COOKIE_ENCRYPTED_SIZE	(COOKIE_COOKIE_SIZE + COOKIE_MAC_SIZE)
 
-#define COOKIE_MAC1_KEY_LABEL	"mac1----"
-#define COOKIE_COOKIE_KEY_LABEL	"cookie--"
-#define COOKIE_SECRET_MAX_AGE	120
-#define COOKIE_SECRET_LATENCY	5
-
-/* Constants for initiation rate limiting */
-#define RATELIMIT_SIZE		(1 << 13)
-#define RATELIMIT_SIZE_MAX	(RATELIMIT_SIZE * 8)
-#define INITIATIONS_PER_SECOND	20
-#define INITIATIONS_BURSTABLE	5
-#define INITIATION_COST		(SBT_1S / INITIATIONS_PER_SECOND)
-#define TOKEN_MAX		(INITIATION_COST * INITIATIONS_BURSTABLE)
-#define ELEMENT_TIMEOUT		1
-#define IPV4_MASK_SIZE		4 /* Use all 4 bytes of IPv4 address */
-#define IPV6_MASK_SIZE		8 /* Use top 8 bytes (/64) of IPv6 address */
-
 struct cookie_macs {
 	uint8_t	mac1[COOKIE_MAC_SIZE];
 	uint8_t	mac2[COOKIE_MAC_SIZE];
 };
 
-struct ratelimit_entry {
-	LIST_ENTRY(ratelimit_entry)	 r_entry;
-	sa_family_t			 r_af;
-	union {
-		struct in_addr		 r_in;
-#ifdef INET6
-		struct in6_addr		 r_in6;
-#endif
-	};
-	sbintime_t			 r_last_time;	/* sbinuptime */
-	uint64_t			 r_tokens;
-};
-
-struct ratelimit {
-	uint8_t				 rl_secret[SIPHASH_KEY_LENGTH];
-	uma_zone_t			 rl_zone;
-
-	struct rwlock			 rl_lock;
-	struct callout			 rl_gc;
-	LIST_HEAD(, ratelimit_entry)	*rl_table;
-	u_long				 rl_table_mask;
-	size_t				 rl_table_num;
-};
-
 struct cookie_maker {
 	uint8_t		cp_mac1_key[COOKIE_KEY_SIZE];
 	uint8_t		cp_cookie_key[COOKIE_KEY_SIZE];
@@ -80,28 +40,24 @@ struct cookie_maker {
 };
 
 struct cookie_checker {
-	struct ratelimit	cc_ratelimit_v4;
-#ifdef INET6
-	struct ratelimit	cc_ratelimit_v6;
-#endif
-
-	struct rwlock		cc_key_lock;
-	uint8_t			cc_mac1_key[COOKIE_KEY_SIZE];
-	uint8_t			cc_cookie_key[COOKIE_KEY_SIZE];
+	struct rwlock	cc_key_lock;
+	uint8_t		cc_mac1_key[COOKIE_KEY_SIZE];
+	uint8_t		cc_cookie_key[COOKIE_KEY_SIZE];
 
-	struct rwlock		cc_secret_lock;
-	sbintime_t		cc_secret_birthdate;	/* sbinuptime */
-	uint8_t			cc_secret[COOKIE_SECRET_SIZE];
+	struct rwlock	cc_secret_lock;
+	sbintime_t	cc_secret_birthdate;	/* sbinuptime */
+	uint8_t		cc_secret[COOKIE_SECRET_SIZE];
 };
 
-void	cookie_maker_init(struct cookie_maker *, const uint8_t[COOKIE_INPUT_SIZE]);
-int	cookie_checker_init(struct cookie_checker *, uma_zone_t);
+int	cookie_init(void);
+void	cookie_deinit(void);
+void	cookie_checker_init(struct cookie_checker *);
 void	cookie_checker_update(struct cookie_checker *,
 	    const uint8_t[COOKIE_INPUT_SIZE]);
-void	cookie_checker_deinit(struct cookie_checker *);
 void	cookie_checker_create_payload(struct cookie_checker *,
 	    struct cookie_macs *cm, uint8_t[COOKIE_NONCE_SIZE],
 	    uint8_t [COOKIE_ENCRYPTED_SIZE], struct sockaddr *);
+void	cookie_maker_init(struct cookie_maker *, const uint8_t[COOKIE_INPUT_SIZE]);
 int	cookie_maker_consume_payload(struct cookie_maker *,
 	    uint8_t[COOKIE_NONCE_SIZE], uint8_t[COOKIE_ENCRYPTED_SIZE]);
 void	cookie_maker_mac(struct cookie_maker *, struct cookie_macs *,