Implement peer limit constraints

6 files changed, 79 insertions, 13 deletions

diff --git a/src/config.l b/src/config.l
index e8d5584..afe37dd 100644
--- a/src/config.l
+++ b/src/config.l
@@ -101,6 +101,7 @@ crypto		{ TOKEN(TOK_CRYPTO); }
 use		{ TOKEN(TOK_USE); }
 default		{ TOKEN(TOK_DEFAULT); }
 group		{ TOKEN(TOK_GROUP); }
+limit		{ TOKEN(TOK_LIMIT); }
 
 [0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3} {
 			UPDATE_LOCATION;
diff --git a/src/config.y b/src/config.y
index e3c8f82..5e3e837 100644
--- a/src/config.y
+++ b/src/config.y
@@ -100,6 +100,7 @@
 %token TOK_USE
 %token TOK_DEFAULT
 %token TOK_GROUP
+%token TOK_LIMIT
 
 %token <addr4> TOK_ADDR4
 %token <addr6> TOK_ADDR6
@@ -152,6 +153,7 @@ statement:	TOK_LOG log ';'
 	|	TOK_ON TOK_DISESTABLISH on_disestablish ';'
 	|	TOK_PEER peer '{' peer_conf '}'
 	|	TOK_PEER TOK_GROUP peer_group '{' config '}' peer_group_after
+	|	TOK_PEER TOK_LIMIT peer_limit ';'
 	|	TOK_FORWARD forward ';'
 	|	TOK_INCLUDE include ';'
 	;
@@ -376,6 +378,11 @@ peer_group_after:
 		}
 	;
 
+peer_limit:	TOK_INTEGER {
+			conf->peer_group->max_connections = $1;
+		}
+	;
+
 forward:	boolean		{ conf->forward = $1; }
 	;
 
diff --git a/src/fastd.c b/src/fastd.c
index d81f949..9e9b225 100644
--- a/src/fastd.c
+++ b/src/fastd.c
@@ -619,19 +619,21 @@ static inline void update_time(fastd_context *ctx) {
 }
 
 static inline void send_handshake(fastd_context *ctx, fastd_peer *peer) {
-	if (!fastd_peer_is_established(peer))
-		fastd_peer_reset_socket(ctx, peer);
+	if (fastd_peer_may_connect(ctx, peer)) {
+		if (!fastd_peer_is_established(peer))
+			fastd_peer_reset_socket(ctx, peer);
 
-	if (peer->sock) {
-		if (timespec_diff(&ctx->now, &peer->last_handshake) < ctx->conf->min_handshake_interval*1000
-		    && fastd_peer_address_equal(&peer->address, &peer->last_handshake_address)) {
-			pr_debug(ctx, "not sending a handshake to %P as we sent one a short time ago", peer);
-		}
-		else {
-			pr_debug(ctx, "sending handshake to %P...", peer);
-			peer->last_handshake = ctx->now;
-			peer->last_handshake_address = peer->address;
-			ctx->conf->protocol->handshake_init(ctx, peer->sock, &peer->address, peer->config);
+		if (peer->sock) {
+			if (timespec_diff(&ctx->now, &peer->last_handshake) < ctx->conf->min_handshake_interval*1000
+			    && fastd_peer_address_equal(&peer->address, &peer->last_handshake_address)) {
+				pr_debug(ctx, "not sending a handshake to %P as we sent one a short time ago", peer);
+			}
+			else {
+				pr_debug(ctx, "sending handshake to %P...", peer);
+				peer->last_handshake = ctx->now;
+				peer->last_handshake_address = peer->address;
+				ctx->conf->protocol->handshake_init(ctx, peer->sock, &peer->address, peer->config);
+			}
 		}
 	}
 
@@ -748,6 +750,11 @@ static void handle_socket(fastd_context *ctx, fastd_socket *sock) {
 	}
 
 	if (peer) {
+		if (!fastd_peer_may_connect(ctx, peer)) {
+			fastd_buffer_free(buffer);
+			return;
+		}
+
 		switch (*packet_type) {
 		case PACKET_DATA:
 			ctx->conf->protocol->handle_recv(ctx, peer, buffer);
diff --git a/src/peer.c b/src/peer.c
index 117d51a..a735f5e 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -404,6 +404,49 @@ static inline fastd_peer_group* find_peer_group(fastd_peer_group *group, const f
 	return NULL;
 }
 
+static inline bool is_group_in(fastd_peer_group *group1, fastd_peer_group *group2) {
+	while (group1) {
+		if (group1 == group2)
+			return true;
+
+		group1 = group1->parent;
+	}
+
+	return false;
+}
+
+static bool is_peer_in_group(fastd_peer *peer, fastd_peer_group *group) {
+	return is_group_in(peer->group, group);
+}
+
+static inline unsigned count_established_group_peers(fastd_context *ctx, fastd_peer_group *group) {
+	unsigned ret = 0;
+	fastd_peer *peer;
+	for (peer = ctx->peers; peer; peer = peer->next) {
+		if (fastd_peer_is_established(peer) && is_peer_in_group(peer, group))
+			ret++;
+	}
+
+	return ret;
+}
+
+bool fastd_peer_may_connect(fastd_context *ctx, fastd_peer *peer) {
+	if (fastd_peer_is_established(peer))
+		return true;
+
+	fastd_peer_group *group;
+
+	for (group = peer->group; group; group = group->parent) {
+		if (!group->conf->max_connections)
+			continue;
+
+		if (count_established_group_peers(ctx, group) >= group->conf->max_connections)
+			return false;
+	}
+
+	return true;
+}
+
 fastd_peer* fastd_peer_add(fastd_context *ctx, fastd_peer_config *peer_conf) {
 	fastd_peer *peer = malloc(sizeof(fastd_peer));
 
@@ -427,6 +470,8 @@ void fastd_peer_set_established(fastd_context *ctx, fastd_peer *peer) {
 		on_establish(ctx, peer);
 		pr_info(ctx, "connection with %P established.", peer);
 	}
+
+	return;
 }
 
 const fastd_eth_addr* fastd_get_source_address(const fastd_context *ctx, fastd_buffer buffer) {
diff --git a/src/peer.h b/src/peer.h
index 1e2cad0..62ca31d 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -90,6 +90,7 @@ void fastd_peer_reset(fastd_context *ctx, fastd_peer *peer);
 void fastd_peer_delete(fastd_context *ctx, fastd_peer *peer);
 fastd_peer* fastd_peer_add(fastd_context *ctx, fastd_peer_config *conf);
 void fastd_peer_set_established(fastd_context *ctx, fastd_peer *peer);
+bool fastd_peer_may_connect(fastd_context *ctx, fastd_peer *peer);
 bool fastd_peer_claim_address(fastd_context *ctx, fastd_peer *peer, fastd_socket *sock, const fastd_peer_address *addr);
 void fastd_peer_reset_socket(fastd_context *ctx, fastd_peer *peer);
 
diff --git a/src/protocol_ec25519_fhmqvc.c b/src/protocol_ec25519_fhmqvc.c
index f2582f2..6ebb771 100644
--- a/src/protocol_ec25519_fhmqvc.c
+++ b/src/protocol_ec25519_fhmqvc.c
@@ -559,6 +559,11 @@ static void protocol_handshake_handle(fastd_context *ctx, fastd_socket *sock, co
 
 	fastd_peer *peer = get_peer(ctx, peer_conf);
 
+	if (!fastd_peer_may_connect(ctx, peer)) {
+		pr_debug(ctx, "ignoring handshake from %P[%I] because of local constraints", peer, address);
+		return;
+	}
+
 	if (backoff(ctx, peer)) {
 		pr_debug(ctx, "received repeated handshakes from %P[%I], ignoring", peer, address);
 		return;
@@ -726,7 +731,7 @@ static void session_send(fastd_context *ctx, fastd_peer *peer, fastd_buffer buff
 }
 
 static void protocol_send(fastd_context *ctx, fastd_peer *peer, fastd_buffer buffer) {
-	if (!peer->protocol_state || !is_session_valid(ctx, &peer->protocol_state->session)) {
+	if (!peer->protocol_state || !fastd_peer_is_established(peer) || !is_session_valid(ctx, &peer->protocol_state->session)) {
 		fastd_buffer_free(buffer);
 		return;
 	}