From 6dcf16386f710304af2a00706f24f683305b8825 Mon Sep 17 00:00:00 2001
From: Linus Nordberg <linus@nordberg.se>
Date: Sun, 29 Sep 2019 22:08:29 +0200
Subject: Remove allowed-ips from peers as leases expire

---
 lease.c             | 20 ++++++++++++++++++-
 lease.h             |  4 +++-
 wg-dynamic-server.c | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/lease.c b/lease.c
index 8a48a27..5ea4f6f 100644
--- a/lease.c
+++ b/lease.c
@@ -203,8 +203,9 @@ bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime)
 	return false;
 }
 
-int leases_refresh()
+int leases_refresh(void (*update_cb)(wg_key *, int))
 {
+	wg_key updates[WG_DYNAMIC_LEASE_CHUNKSIZE] = { 0 };
 	time_t cur_time = get_monotonic_time();
 
 	if (cur_time < gexpires)
@@ -212,6 +213,7 @@ int leases_refresh()
 
 	gexpires = TIME_T_MAX;
 
+	int i = 0;
 	for (khint_t k = kh_begin(leases_ht); k != kh_end(leases_ht); ++k) {
 		if (!kh_exist(leases_ht, k))
 			continue;
@@ -228,6 +230,19 @@ int leases_refresh()
 				if (!IN6_IS_ADDR_UNSPECIFIED(ipv6))
 					ipp_del_v6(&pool, ipv6, 128);
 
+				memcpy(updates[i], kh_key(leases_ht, k), sizeof(wg_key));
+				{
+					wg_key_b64_string pubkey_asc;
+					wg_key_to_base64(pubkey_asc, updates[i]);
+					debug("Peer losing its lease: %s\n", pubkey_asc);
+				}
+				i++;
+				if (i == WG_DYNAMIC_LEASE_CHUNKSIZE) {
+					update_cb(updates, i);
+					i = 0;
+					memset(updates, 0, sizeof updates);
+				}
+
 				tmp = *pp;
 				*pp = (*pp)->next;
 				free(tmp);
@@ -239,6 +254,9 @@ int leases_refresh()
 			}
 		}
 
+		if (i)
+			update_cb(updates, i);
+
 		if (!kh_val(leases_ht, k)) {
 			free((char *)kh_key(leases_ht, k));
 			kh_del(leaseht, leases_ht, k);
diff --git a/lease.h b/lease.h
index 3e1402d..23852e6 100644
--- a/lease.h
+++ b/lease.h
@@ -13,6 +13,8 @@
 #include "common.h"
 #include "netlink.h"
 
+#define WG_DYNAMIC_LEASE_CHUNKSIZE 256
+
 struct wg_dynamic_lease {
 	time_t start_real;
 	time_t start_mono;
@@ -54,7 +56,7 @@ bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime);
 /* Refreshes all leases, meaning expired ones will be removed. Returns the
  * amount of seconds until the next lease will expire, or at most INT_MAX/1000.
  */
-int leases_refresh();
+int leases_refresh(void (*update_cb)(wg_key *, int));
 
 /*
  * Updates all pools with information from the mnl socket nlsock.
diff --git a/wg-dynamic-server.c b/wg-dynamic-server.c
index fca1dfc..d1229cc 100644
--- a/wg-dynamic-server.c
+++ b/wg-dynamic-server.c
@@ -294,6 +294,59 @@ static void add_allowed_ips(wg_key pubkey, struct in_addr *ipv4,
 		fatal("wg_set_device()");
 }
 
+/* TODO: have UPDATES contain {wg_key, ip4 and ip6} and remove only matching addrs */
+static void update_allowed_ips(wg_key *updates, int nupdates)
+{
+	wg_device newdev = { 0 };
+	wg_peer newpeers[WG_DYNAMIC_LEASE_CHUNKSIZE] = { 0 },
+		**nextpp = &newdev.first_peer;
+	wg_allowedip newallowedips[WG_DYNAMIC_LEASE_CHUNKSIZE] = { 0 };
+
+	int newpeers_idx = 0;
+	wg_peer *peer;
+	wg_for_each_peer (device, peer) {
+		for (int i = 0; i < nupdates; i++) {
+			if (!memcmp(peer->public_key, updates[i],
+				    sizeof(wg_key))) {
+				wg_peer *pp = &newpeers[newpeers_idx];
+				pp->flags |= WGPEER_REPLACE_ALLOWEDIPS;
+				memcpy(pp->public_key, peer->public_key,
+				       sizeof(wg_key));
+
+				/* Whacking all addrs except the first (!) link-local /128 . */
+				wg_allowedip *allowedip;
+				wg_for_each_allowedip (peer, allowedip) {
+					if (allowedip->family == AF_INET6 &&
+					    allowedip->cidr == 128 &&
+					    IN6_IS_ADDR_LINKLOCAL(
+						    &allowedip->ip6)) {
+						wg_allowedip *aip =
+							&newallowedips
+								[newpeers_idx];
+						aip->family = AF_INET6;
+						memcpy(&aip->ip6,
+						       &allowedip->ip6,
+						       sizeof(struct in6_addr));
+						aip->cidr = 128;
+						pp->first_allowedip = aip;
+						break;
+					}
+				}
+				newpeers_idx++;
+				*nextpp = pp;
+				nextpp = &pp->next_peer;
+				break; /* Assuming no duplicated pubkeys in updates. */
+			}
+		}
+	}
+
+	if (newpeers_idx) {
+		strcpy(newdev.name, wg_interface);
+		if (wg_set_device(&newdev))
+			fatal("wg_set_device()");
+	}
+}
+
 static int response_request_ip(struct wg_dynamic_attr *cur, wg_key pubkey,
 			       struct wg_dynamic_lease **lease)
 {
@@ -546,7 +599,7 @@ static void poll_loop()
 		fatal("epoll_ctl()");
 
 	while (1) {
-		time_t next = leases_refresh() * 1000;
+		time_t next = leases_refresh(update_allowed_ips) * 1000;
 		int nfds = epoll_wait(epollfd, events, MAX_CONNECTIONS, next);
 		if (nfds == -1) {
 			if (errno == EINTR)
-- 
cgit v1.2.3-59-g8ed1b