wg_cookie: hash vnet into ratelimiter entry

IPs mean different things per-vnet.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

5 files changed, 46 insertions, 62 deletions

diff --git a/TODO.md b/TODO.md
index 5685531..1f89520 100644
--- a/TODO.md
+++ b/TODO.md
@@ -11,10 +11,7 @@
 - Make sure noise state machine is correct.
 - Investigate whether the allowed ips lookup structure needs reference
   counting.
-- Handle failures of `rn_inithead` and remember to call `rn_detachhead`
-  somewhere during cleanup.
 - Stop using `M_WAITOK` and use `M_NOWAIT` instead.
-- Hash in jail/fib pointer to ratelimiter.
 
 ### Crypto TODO
 
diff --git a/src/if_wg.c b/src/if_wg.c
index b678f68..8680337 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -1288,7 +1288,8 @@ wg_handshake(struct wg_softc *sc, struct wg_packet *pkt)
 
 		res = cookie_checker_validate_macs(&sc->sc_cookie, &init->m,
 				init, sizeof(*init) - sizeof(init->m),
-				underload, &e->e_remote.r_sa);
+				underload, &e->e_remote.r_sa,
+				sc->sc_ifp->if_vnet);
 
 		if (res == EINVAL) {
 			DPRINTF(sc, "Invalid initiation MAC\n");
@@ -1321,7 +1322,8 @@ wg_handshake(struct wg_softc *sc, struct wg_packet *pkt)
 
 		res = cookie_checker_validate_macs(&sc->sc_cookie, &resp->m,
 				resp, sizeof(*resp) - sizeof(resp->m),
-				underload, &e->e_remote.r_sa);
+				underload, &e->e_remote.r_sa,
+				sc->sc_ifp->if_vnet);
 
 		if (res == EINVAL) {
 			DPRINTF(sc, "Invalid response MAC\n");
diff --git a/src/selftest/cookie.c b/src/selftest/cookie.c
index 5205ea2..d5778b7 100644
--- a/src/selftest/cookie.c
+++ b/src/selftest/cookie.c
@@ -55,7 +55,7 @@ cookie_ratelimit_timings_test(void)
 		sin.sin_addr.s_addr = 0x01020304;
 		sin.sin_port = arc4random();
 
-		if (ratelimit_allow(&rl, sintosa(&sin)) != rl_expected[i].result)
+		if (ratelimit_allow(&rl, sintosa(&sin), NULL) != rl_expected[i].result)
 			T_FAILED_ITER("malicious v4");
 
 		/* The second ratelimit_allow is to test that an arbitrary
@@ -63,7 +63,7 @@ cookie_ratelimit_timings_test(void)
 		sin.sin_addr.s_addr += i + 1;
 		sin.sin_port = arc4random();
 
-		if (ratelimit_allow(&rl, sintosa(&sin)) != 0)
+		if (ratelimit_allow(&rl, sintosa(&sin), NULL) != 0)
 			T_FAILED_ITER("non-malicious v4");
 
 #ifdef INET6
@@ -77,7 +77,7 @@ cookie_ratelimit_timings_test(void)
 		sin6.sin6_addr.s6_addr32[3] = i;
 		sin6.sin6_port = arc4random();
 
-		if (ratelimit_allow(&rl, sin6tosa(&sin6)) != rl_expected[i].result)
+		if (ratelimit_allow(&rl, sin6tosa(&sin6), NULL) != rl_expected[i].result)
 			T_FAILED_ITER("malicious v6");
 
 		/* Again, test that an address different to above is still
@@ -85,7 +85,7 @@ cookie_ratelimit_timings_test(void)
 		sin6.sin6_addr.s6_addr32[0] += i + 1;
 		sin6.sin6_port = arc4random();
 
-		if (ratelimit_allow(&rl, sintosa(&sin)) != 0)
+		if (ratelimit_allow(&rl, sintosa(&sin), NULL) != 0)
 			T_FAILED_ITER("non-malicious v6");
 #endif
 	}
@@ -110,10 +110,10 @@ cookie_ratelimit_capacity_test(void)
 	for (i = 0; i <= RATELIMIT_SIZE_MAX; i++) {
 		sin.sin_addr.s_addr = i;
 		if (i == RATELIMIT_SIZE_MAX) {
-			if (ratelimit_allow(&rl, sintosa(&sin)) != ECONNREFUSED)
+			if (ratelimit_allow(&rl, sintosa(&sin), NULL) != ECONNREFUSED)
 				T_FAILED_ITER("reject");
 		} else {
-			if (ratelimit_allow(&rl, sintosa(&sin)) != 0)
+			if (ratelimit_allow(&rl, sintosa(&sin), NULL) != 0)
 				T_FAILED_ITER("allow");
 		}
 	}
@@ -139,7 +139,7 @@ cookie_ratelimit_gc_test(void)
 
 	for (i = 0; i < RATELIMIT_SIZE_MAX / 2; i++) {
 		sin.sin_addr.s_addr = i;
-		if (ratelimit_allow(&rl, sintosa(&sin)) != 0)
+		if (ratelimit_allow(&rl, sintosa(&sin), NULL) != 0)
 			T_FAILED_ITER("insert");
 	}
 
@@ -150,7 +150,7 @@ cookie_ratelimit_gc_test(void)
 
 	for (i = 0; i < RATELIMIT_SIZE_MAX / 2; i++) {
 		sin.sin_addr.s_addr = i;
-		if (ratelimit_allow(&rl, sintosa(&sin)) != 0)
+		if (ratelimit_allow(&rl, sintosa(&sin), NULL) != 0)
 			T_FAILED_ITER("insert");
 	}
 
@@ -207,7 +207,7 @@ cookie_mac_test(void)
 	for (i = 0; i < sizeof(cm.mac1); i++) {
 		cm.mac1[i] = ~cm.mac1[i];
 		if (cookie_checker_validate_macs(&checker, &cm, message,
-		    MESSAGE_LEN, 0, sintosa(&sin)) != EINVAL)
+		    MESSAGE_LEN, 0, sintosa(&sin), NULL) != EINVAL)
 			T_FAILED("validate_macs_noload_munge");
 		cm.mac1[i] = ~cm.mac1[i];
 	}
@@ -222,12 +222,12 @@ cookie_mac_test(void)
 
 	/* Check we can successfully validate the MAC */
 	if (cookie_checker_validate_macs(&checker, &cm, message,
-	    MESSAGE_LEN, 0, sintosa(&sin)) != 0)
+	    MESSAGE_LEN, 0, sintosa(&sin), NULL) != 0)
 		T_FAILED("validate_macs_noload_normal");
 
 	/* Check we get a EAGAIN if no mac2 and under load */
 	if (cookie_checker_validate_macs(&checker, &cm, message,
-	    MESSAGE_LEN, 1, sintosa(&sin)) != EAGAIN)
+	    MESSAGE_LEN, 1, sintosa(&sin), NULL) != EAGAIN)
 		T_FAILED("validate_macs_load_normal");
 
 	/* Simulate a cookie message */
@@ -261,19 +261,19 @@ cookie_mac_test(void)
 
 	/* Check we get OK if mac2 and under load */
 	if (cookie_checker_validate_macs(&checker, &cm, message,
-	    MESSAGE_LEN, 1, sintosa(&sin)) != 0)
+	    MESSAGE_LEN, 1, sintosa(&sin), NULL) != 0)
 		T_FAILED("validate_macs_load_normal_mac2");
 
 	sin.sin_addr.s_addr = ~sin.sin_addr.s_addr;
 	/* Check we get EAGAIN if we munge the source IP */
 	if (cookie_checker_validate_macs(&checker, &cm, message,
-	    MESSAGE_LEN, 1, sintosa(&sin)) != EAGAIN)
+	    MESSAGE_LEN, 1, sintosa(&sin), NULL) != EAGAIN)
 		T_FAILED("validate_macs_load_spoofip_mac2");
 	sin.sin_addr.s_addr = ~sin.sin_addr.s_addr;
 
 	/* Check we get OK if mac2 and under load */
 	if (cookie_checker_validate_macs(&checker, &cm, message,
-	    MESSAGE_LEN, 1, sintosa(&sin)) != 0)
+	    MESSAGE_LEN, 1, sintosa(&sin), NULL) != 0)
 		T_FAILED("validate_macs_load_normal_mac2_retry");
 
 	T_PASSED;
diff --git a/src/wg_cookie.c b/src/wg_cookie.c
index e68d662..58d0d8d 100644
--- a/src/wg_cookie.c
+++ b/src/wg_cookie.c
@@ -31,22 +31,20 @@
 #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_key {
+	struct vnet *vnet;
+	uint8_t ip[IPV6_MASK_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
-	};
+	struct ratelimit_key		r_key;
 	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[RATELIMIT_SIZE];
@@ -67,7 +65,7 @@ static void	ratelimit_deinit(struct ratelimit *);
 static void	ratelimit_gc_callout(void *);
 static void	ratelimit_gc_schedule(struct ratelimit *);
 static void	ratelimit_gc(struct ratelimit *, bool);
-static int	ratelimit_allow(struct ratelimit *, struct sockaddr *);
+static int	ratelimit_allow(struct ratelimit *, struct sockaddr *, struct vnet *);
 static uint64_t siphash13(const uint8_t [SIPHASH_KEY_LENGTH], const void *, size_t);
 
 static struct ratelimit ratelimit_v4;
@@ -86,9 +84,9 @@ cookie_init(void)
 
 	ratelimit_init(&ratelimit_v4);
 #ifdef INET6
-	ratelimit_init(&ratelimit_v6)
+	ratelimit_init(&ratelimit_v6);
 #endif
-	return 0;
+	return (0);
 }
 
 void
@@ -207,7 +205,7 @@ cookie_maker_mac(struct cookie_maker *cm, struct cookie_macs *macs, void *buf,
 
 int
 cookie_checker_validate_macs(struct cookie_checker *cc, struct cookie_macs *macs,
-    void *buf, size_t len, bool check_cookie, struct sockaddr *sa)
+    void *buf, size_t len, bool check_cookie, struct sockaddr *sa, struct vnet *vnet)
 {
 	struct cookie_macs our_macs;
 	uint8_t cookie[COOKIE_COOKIE_SIZE];
@@ -234,10 +232,10 @@ cookie_checker_validate_macs(struct cookie_checker *cc, struct cookie_macs *macs
 		 * implying there is no ratelimiting, or we should ratelimit
 		 * (refuse) respectively. */
 		if (sa->sa_family == AF_INET)
-			return ratelimit_allow(&ratelimit_v4, sa);
+			return ratelimit_allow(&ratelimit_v4, sa, vnet);
 #ifdef INET6
 		else if (sa->sa_family == AF_INET6)
-			return ratelimit_allow(&ratelimit_v6, sa);
+			return ratelimit_allow(&ratelimit_v6, sa, vnet);
 #endif
 		else
 			return EAFNOSUPPORT;
@@ -391,40 +389,33 @@ ratelimit_gc(struct ratelimit *rl, bool force)
 }
 
 static int
-ratelimit_allow(struct ratelimit *rl, struct sockaddr *sa)
+ratelimit_allow(struct ratelimit *rl, struct sockaddr *sa, struct vnet *vnet)
 {
-	uint64_t key, tokens;
+	uint64_t bucket, tokens;
 	sbintime_t diff, now;
 	struct ratelimit_entry *r;
 	int ret = ECONNREFUSED;
+	struct ratelimit_key key = { .vnet = vnet };
+	size_t len = sizeof(key);
 
-	if (sa->sa_family == AF_INET)
-		key = siphash13(rl->rl_secret, &satosin(sa)->sin_addr,
-				IPV4_MASK_SIZE);
+	if (sa->sa_family == AF_INET) {
+		memcpy(key.ip, &satosin(sa)->sin_addr, IPV4_MASK_SIZE);
+		len -= IPV6_MASK_SIZE - IPV4_MASK_SIZE;
+	}
 #ifdef INET6
 	else if (sa->sa_family == AF_INET6)
-		key = siphash13(rl->rl_secret, &satosin6(sa)->sin6_addr,
-				IPV6_MASK_SIZE);
+		memcpy(key.ip, &satosin6(sa)->sin6_addr, IPV6_MASK_SIZE);
 #endif
 	else
 		return ret;
 
+	bucket = siphash13(rl->rl_secret, &key, len) & RATELIMIT_MASK;
 	rw_wlock(&rl->rl_lock);
 
-	LIST_FOREACH(r, &rl->rl_table[key & RATELIMIT_MASK], r_entry) {
-		if (r->r_af != sa->sa_family)
+	LIST_FOREACH(r, &rl->rl_table[bucket], r_entry) {
+		if (bcmp(&r->r_key, &key, len) != 0)
 			continue;
 
-		if (r->r_af == AF_INET && bcmp(&r->r_in,
-		    &satosin(sa)->sin_addr, IPV4_MASK_SIZE) != 0)
-			continue;
-
-#ifdef INET6
-		if (r->r_af == AF_INET6 && bcmp(&r->r_in6,
-		    &satosin6(sa)->sin6_addr, IPV6_MASK_SIZE) != 0)
-			continue;
-#endif
-
 		/* If we get to here, we've found an entry for the endpoint.
 		 * We apply standard token bucket, by calculating the time
 		 * lapsed since our last_time, adding that, ensuring that we
@@ -462,15 +453,8 @@ ratelimit_allow(struct ratelimit *rl, struct sockaddr *sa)
 	rl->rl_table_num++;
 
 	/* Insert entry into the hashtable and ensure it's initialised */
-	LIST_INSERT_HEAD(&rl->rl_table[key & RATELIMIT_MASK], r, r_entry);
-	r->r_af = sa->sa_family;
-	if (r->r_af == AF_INET)
-		memcpy(&r->r_in, &satosin(sa)->sin_addr, IPV4_MASK_SIZE);
-#ifdef INET6
-	else if (r->r_af == AF_INET6)
-		memcpy(&r->r_in6, &satosin6(sa)->sin6_addr, IPV6_MASK_SIZE);
-#endif
-
+	LIST_INSERT_HEAD(&rl->rl_table[bucket], r, r_entry);
+	r->r_key = key;
 	r->r_last_time = getsbinuptime();
 	r->r_tokens = TOKEN_MAX - INITIATION_COST;
 
diff --git a/src/wg_cookie.h b/src/wg_cookie.h
index a2e20f2..099cda6 100644
--- a/src/wg_cookie.h
+++ b/src/wg_cookie.h
@@ -64,7 +64,8 @@ int	cookie_maker_consume_payload(struct cookie_maker *,
 void	cookie_maker_mac(struct cookie_maker *, struct cookie_macs *,
 	    void *, size_t);
 int	cookie_checker_validate_macs(struct cookie_checker *,
-	    struct cookie_macs *, void *, size_t, bool, struct sockaddr *);
+	    struct cookie_macs *, void *, size_t, bool, struct sockaddr *,
+	    struct vnet *);
 
 #ifdef SELFTESTS
 void	cookie_selftest(void);