From 4138d26ca9b711aa63e8b474d6c183f91219b3a4 Mon Sep 17 00:00:00 2001
From: Michael Tremer <michael.tremer@ipfire.org>
Date: Thu, 18 Feb 2021 16:51:13 +0000
Subject: [PATCH] server: Support "--close"

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---
 man/fireperf.txt | 14 +++++++-------
 src/client.c     | 26 ++++++++++----------------
 src/server.c     |  9 +++++++++
 3 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/man/fireperf.txt b/man/fireperf.txt
index fa100f0..c519361 100644
--- a/man/fireperf.txt
+++ b/man/fireperf.txt
@@ -45,6 +45,13 @@ a network between a fireperf client and fireperf server.
 -V::
 	Prints the version of this program
 
+--close::
+-x::
+	If set, all connections will be closed as soon as they have been established.
+	No data will be sent.
+	+
+	This is useful to test how many TCP handshakes per second are possible.
+
 --zero::
 -z::
 	Instead of sending random data, this option will set the client to send packets
@@ -87,13 +94,6 @@ a network between a fireperf client and fireperf server.
 	If set, the client will automatically terminate itself after T seconds.
 	Otherwise it will run for forever.
 
---close::
--x::
-	If set, all connections will be closed as soon as they have been established.
-	No data will be sent.
-	+
-	This is useful to test how many TCP handshakes per second are possible.
-
 
 == BUGS
 Please report all bugs to the bugtracker at https://bugzilla.ipfire.org/
diff --git a/src/client.c b/src/client.c
index 48b377a..aae4b8e 100644
--- a/src/client.c
+++ b/src/client.c
@@ -135,21 +135,6 @@ ERROR:
 	return -1;
 }
 
-static int handle_connection_ready(struct fireperf_config* conf,
-		struct fireperf_stats* stats, int fd) {
-	// Are we supposed to close this connection straight away?
-	if (conf->close) {
-		DEBUG(conf, "Closing connection %d\n", fd);
-		close(fd);
-
-		stats->open_connections--;
-
-		return 0;
-	}
-
-	return handle_connection_send(conf, stats, fd);
-}
-
 int fireperf_client(struct fireperf_config* conf, struct fireperf_stats* stats,
 		int epollfd, int timerfd) {
 	DEBUG(conf, "Launching " PACKAGE_NAME " in client mode\n");
@@ -244,6 +229,15 @@ int fireperf_client(struct fireperf_config* conf, struct fireperf_stats* stats,
 
 					stats->open_connections--;
 				} else {
+					// Close connections immediately when -x is set
+					if (conf->close) {
+						DEBUG(conf, "Closing connection %d\n", fd);
+						close(fd);
+
+						stats->open_connections--;
+						continue;
+					}
+
 					// Handle incoming data
 					if (events[i].events & EPOLLIN) {
 						r = handle_connection_recv(conf, stats, fd);
@@ -253,7 +247,7 @@ int fireperf_client(struct fireperf_config* conf, struct fireperf_stats* stats,
 
 					// Handle outgoing data
 					if (events[i].events & EPOLLOUT) {
-						r = handle_connection_ready(conf, stats, fd);
+						r = handle_connection_send(conf, stats, fd);
 						if (r)
 							goto ERROR;
 					}
diff --git a/src/server.c b/src/server.c
index 52a09a6..6b21f5d 100644
--- a/src/server.c
+++ b/src/server.c
@@ -220,6 +220,15 @@ int fireperf_server(struct fireperf_config* conf, struct fireperf_stats* stats,
 					continue;
 				}
 
+				// Close connections immediately when -x is set
+				if (conf->close) {
+					DEBUG(conf, "Closing connection %d\n", fd);
+					close(fd);
+
+					stats->open_connections--;
+					continue;
+				}
+
 				// Handle incoming data
 				if (events[i].events & EPOLLIN) {
 					r = handle_connection_recv(conf, stats, fd);
-- 
2.47.2