Implement lease expiration

8 files changed, 120 insertions, 36 deletions

diff --git a/common.c b/common.c
index 09a02e2..28f0ba0 100644
--- a/common.c
+++ b/common.c
@@ -12,7 +12,6 @@
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
-#include <sys/param.h>
 
 #include <arpa/inet.h>
 #include <libmnl/libmnl.h>
diff --git a/dbg.h b/dbg.h
index 2d32033..3432fd3 100644
--- a/dbg.h
+++ b/dbg.h
@@ -66,4 +66,6 @@ extern int DBG_LVL;
 /* A neat macro that silences unused parameter warnings compiler independant */
 #define UNUSED(x) (void)(x)
 
+#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y))
+
 #endif
diff --git a/lease.c b/lease.c
index c13a190..b105bc5 100644
--- a/lease.c
+++ b/lease.c
@@ -15,8 +15,10 @@
 #include "radix-trie.h"
 #include "random.h"
 
+#define TIME_T_MAX (((time_t)1 << (sizeof(time_t) * CHAR_BIT - 2)) - 1) * 2 + 1
+
 static struct ip_pool pools;
-static time_t gexpires;
+static time_t gexpires = TIME_T_MAX;
 
 KHASH_MAP_INIT_WGKEY(leaseht, struct wg_dynamic_lease *)
 khash_t(leaseht) * leases_ht;
@@ -75,8 +77,7 @@ void leases_free()
 }
 
 struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
-				   struct in_addr *ipv4, struct in6_addr *ipv6,
-				   time_t *expires)
+				   struct in_addr *ipv4, struct in6_addr *ipv6)
 {
 	struct wg_dynamic_lease *lease, *parent;
 	uint64_t index_low;
@@ -138,8 +139,13 @@ struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
 			ipp_addnth_v6(&pools, &lease->ipv6, index_low,
 				      index_high);
 		} else {
-			if (ipp_add_v6(&pools, ipv6, 128))
-				return NULL; /* TODO: free ipv4 addr */
+			if (ipp_add_v6(&pools, ipv6, 128)) {
+				if (!ipv4 || ipv4->s_addr)
+					ipp_del_v4(&pools, ipv4, 32);
+
+				return NULL;
+			}
+
 			memcpy(&lease->ipv6, ipv6, sizeof *ipv6);
 		}
 	}
@@ -154,7 +160,7 @@ struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
 
 	k = kh_put(leaseht, leases_ht, pubkey, &ret);
 	if (ret < 0) {
-		die("kh_put()");
+		fatal("kh_put()");
 	} else if (ret == 0) {
 		parent = kh_value(leases_ht, k);
 		while (parent->next)
@@ -165,10 +171,8 @@ struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
 		kh_value(leases_ht, k) = lease;
 	}
 
-	if (lease->start_mono < gexpires)
-		gexpires = lease->start_mono;
-
-	*expires = gexpires;
+	if (lease->start_mono + lease->leasetime < gexpires)
+		gexpires = lease->start_mono + lease->leasetime;
 
 	/* TODO: add record to file */
 
@@ -185,19 +189,59 @@ struct wg_dynamic_lease *get_leases(wg_key pubkey)
 		return kh_val(leases_ht, k);
 }
 
-bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime,
-		  time_t *expires)
+bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime)
 {
 	UNUSED(lease);
 	UNUSED(leasetime);
-	UNUSED(expires);
 	return false;
 }
 
-time_t leases_refresh()
+int leases_refresh()
 {
-	/* TODO: remove expired leases */
-	return gexpires;
+	time_t cur_time = get_monotonic_time();
+
+	if (cur_time < gexpires)
+		return MIN(INT_MAX / 1000, gexpires - cur_time);
+
+	gexpires = TIME_T_MAX;
+
+	for (khint_t k = kh_begin(leases_ht); k != kh_end(leases_ht); ++k) {
+		if (!kh_exist(leases_ht, k))
+			continue;
+
+		struct wg_dynamic_lease **pp = &kh_val(leases_ht, k), *tmp;
+		while (*pp) {
+			struct in_addr *ipv4 = &(*pp)->ipv4;
+			struct in6_addr *ipv6 = &(*pp)->ipv6;
+			time_t expires = (*pp)->start_mono + (*pp)->leasetime;
+			if (cur_time >= expires) {
+				if (ipv4->s_addr) {
+					ipp_del_v4(&pools, ipv4, 32);
+					++total_ipv4;
+				}
+				if (!IN6_IS_ADDR_UNSPECIFIED(ipv6)) {
+					ipp_del_v6(&pools, ipv6, 128);
+					++totall_ipv6;
+					if (totall_ipv6 == 0)
+						++totalh_ipv6;
+				}
+
+				tmp = *pp;
+				*pp = (*pp)->next;
+				free(tmp);
+			} else {
+				if (expires < gexpires)
+					gexpires = expires;
+
+				pp = &(*pp)->next;
+			}
+		}
+
+		if (!kh_val(leases_ht, k))
+			kh_del(leaseht, leases_ht, k);
+	}
+
+	return MIN(INT_MAX / 1000, gexpires - cur_time);
 }
 
 void leases_update_pools(int fd)
diff --git a/lease.h b/lease.h
index 4052bf9..95a7b58 100644
--- a/lease.h
+++ b/lease.h
@@ -35,12 +35,9 @@ void leases_free();
 /*
  * Creates a new lease and returns a pointer to it, or NULL if either we ran out
  * of assignable IPs or the requested IP is already taken.
- * expires contains the (monotonic) timestamp after which the next lease,
- * possibly the newly created one, will expire.
  */
 struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
-				   struct in_addr *ipv4, struct in6_addr *ipv6,
-				   time_t *expires);
+				   struct in_addr *ipv4, struct in6_addr *ipv6);
 
 /*
  * Returns all leases belonging to pubkey, or NULL if there are none.
@@ -48,17 +45,15 @@ struct wg_dynamic_lease *new_lease(wg_key pubkey, uint32_t leasetime,
 struct wg_dynamic_lease *get_leases(wg_key pubkey);
 
 /*
- * Extend the lease to be leasetime long again. Returns true on error, or false
- * otherwise. expires behaves exactly as in new_lease().
+ * Extend the lease to be leasetime seconds long again. Returns true on error,
+ * or false otherwise.
  */
-bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime,
-		  time_t *expires);
+bool extend_lease(struct wg_dynamic_lease *lease, uint32_t leasetime);
 
-/*
- * Refreshes all leases, meaning expired ones will be removed. Returns the
- * expiration timestamp of the lease that will expire next.
+/* 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.
  */
-time_t leases_refresh();
+int leases_refresh();
 
 /*
  * Updates all pools with information from the netlink file descriptor fd.
diff --git a/radix-trie.c b/radix-trie.c
index 00d9941..f88d5f9 100644
--- a/radix-trie.c
+++ b/radix-trie.c
@@ -16,8 +16,6 @@
 #include "dbg.h"
 #include "radix-trie.h"
 
-#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y))
-
 #ifndef __aligned
 #define __aligned(x) __attribute__((aligned(x)))
 #endif
@@ -116,7 +114,7 @@ static bool prefix_matches(const struct radix_node *node, const uint8_t *key,
 }
 
 #define CHOOSE_NODE(parent, key)                                               \
-	parent->bit[(key[parent->bit_at_a] >> parent->bit_at_b) & 1]
+	(parent)->bit[(key[(parent)->bit_at_a] >> (parent)->bit_at_b) & 1]
 
 static bool node_placement(struct radix_node *trie, const uint8_t *key,
 			   uint8_t cidr, uint8_t bits,
@@ -350,6 +348,34 @@ static int insert_v6(struct radix_node **root, const struct in6_addr *ip,
 	return ret;
 }
 
+static int remove_node(struct radix_node *trie, const uint8_t *key,
+		       uint8_t bits)
+{
+	struct radix_node **node = &trie, **target = NULL;
+
+	while (*node && prefix_matches(*node, key, bits)) {
+		if ((*node)->is_leaf) {
+			target = node;
+			break;
+		}
+
+		if (CHOOSE_NODE(*node, key) == (*node)->bit[0])
+			++((*node)->left);
+		else
+			++((*node)->right);
+
+		node = &CHOOSE_NODE(*node, key);
+	}
+
+	if (!target)
+		return 1; /* key not found in trie */
+
+	*target = NULL;
+	radix_free_nodes(*node);
+
+	return 0;
+}
+
 static int ipp_addpool(struct radix_pool **pool, struct radix_node **root,
 		       uint8_t bits, const uint8_t *key, uint8_t cidr)
 {
@@ -479,6 +505,22 @@ int ipp_add_v6(struct ip_pool *pool, const struct in6_addr *ip, uint8_t cidr)
 	return insert_v6(&pool->ip6_root, ip, cidr);
 }
 
+int ipp_del_v4(struct ip_pool *pool, const struct in_addr *ip, uint8_t cidr)
+{
+	uint8_t key[4] __aligned(__alignof(uint32_t));
+	swap_endian(key, (const uint8_t *)ip, 32);
+
+	return remove_node(pool->ip4_root, key, cidr);
+}
+
+int ipp_del_v6(struct ip_pool *pool, const struct in6_addr *ip, uint8_t cidr)
+{
+	uint8_t key[16] __aligned(__alignof(uint64_t));
+	swap_endian(key, (const uint8_t *)ip, 128);
+
+	return remove_node(pool->ip6_root, key, cidr);
+}
+
 int ipp_addpool_v4(struct ip_pool *pool, const struct in_addr *ip, uint8_t cidr)
 {
 	uint8_t key[4] __aligned(__alignof(uint32_t));
diff --git a/radix-trie.h b/radix-trie.h
index 6ffaccf..9475cdb 100644
--- a/radix-trie.h
+++ b/radix-trie.h
@@ -21,6 +21,9 @@ void ipp_free(struct ip_pool *pool);
 int ipp_add_v4(struct ip_pool *pool, const struct in_addr *ip, uint8_t cidr);
 int ipp_add_v6(struct ip_pool *pool, const struct in6_addr *ip, uint8_t cidr);
 
+int ipp_del_v4(struct ip_pool *pool, const struct in_addr *ip, uint8_t cidr);
+int ipp_del_v6(struct ip_pool *pool, const struct in6_addr *ip, uint8_t cidr);
+
 uint32_t ipp_gettotal_v4(struct ip_pool *pool);
 uint64_t ipp_gettotal_v6(struct ip_pool *pool, uint32_t *high);
 
diff --git a/wg-dynamic-client.c b/wg-dynamic-client.c
index 8dcbd80..f3e3274 100644
--- a/wg-dynamic-client.c
+++ b/wg-dynamic-client.c
@@ -14,7 +14,6 @@
 #include <arpa/inet.h>
 #include <libmnl/libmnl.h>
 #include <linux/rtnetlink.h>
-#include <sys/param.h>
 
 #include "common.h"
 #include "dbg.h"
diff --git a/wg-dynamic-server.c b/wg-dynamic-server.c
index 8a552a8..e952be0 100644
--- a/wg-dynamic-server.c
+++ b/wg-dynamic-server.c
@@ -291,7 +291,6 @@ static void add_allowed_ips(wg_key pubkey, struct in_addr *ipv4,
 static int response_request_ip(struct wg_dynamic_attr *cur, wg_key pubkey,
 			       struct wg_dynamic_lease **lease)
 {
-	time_t expires;
 	struct in_addr *ipv4 = NULL;
 	struct in6_addr *ipv6 = NULL;
 	uint32_t leasetime = WG_DYNAMIC_LEASETIME;
@@ -319,7 +318,7 @@ static int response_request_ip(struct wg_dynamic_attr *cur, wg_key pubkey,
 	if (ipv4 && ipv6 && !ipv4->s_addr && IN6_IS_ADDR_UNSPECIFIED(ipv6))
 		return 2; /* TODO: invalid request */
 
-	*lease = new_lease(pubkey, leasetime, ipv4, ipv6, &expires);
+	*lease = new_lease(pubkey, leasetime, ipv4, ipv6);
 	if (!*lease)
 		return 1; /* TODO: either out of IPs or IP unavailable */
 
@@ -524,7 +523,8 @@ static void poll_loop()
 		fatal("epoll_ctl()");
 
 	while (1) {
-		int nfds = epoll_wait(epollfd, events, MAX_CONNECTIONS, -1);
+		time_t next = leases_refresh() * 1000;
+		int nfds = epoll_wait(epollfd, events, MAX_CONNECTIONS, next);
 		if (nfds == -1) {
 			if (errno == EINTR)
 				continue;