From 1776e7af9b25cfbb1e28eca94d9052563dad3724 Mon Sep 17 00:00:00 2001
From: David Gibson <david@gibson.dropbear.id.au>
Date: Wed, 15 Nov 2023 16:25:32 +1100
Subject: [PATCH] tcp: Use common helper for rebinding inbound and outbound
 ports

tcp_port_rebind() has two cases with almost but not quite identical code.
Simplify things a bit by factoring this out into a single parameterised
helper, tcp_port_do_rebind().

Signed-off-by: David Gibson <david@gibson.dropbear.id.au>
Signed-off-by: Stefano Brivio <sbrivio@redhat.com>
---
 tcp.c | 92 +++++++++++++++++++++++++++++------------------------------
 1 file changed, 45 insertions(+), 47 deletions(-)

diff --git a/tcp.c b/tcp.c
index cfcd40a..cb23204 100644
--- a/tcp.c
+++ b/tcp.c
@@ -3150,6 +3150,49 @@ int tcp_init(struct ctx *c)
 	return 0;
 }
 
+/**
+ * tcp_port_do_rebind() - Rebind ports to match forward maps
+ * @c:		Execution context
+ * @outbound:	True to remap outbound forwards, otherwise inbound
+ *
+ * Must be called in namespace context if @outbound is true.
+ */
+static void tcp_port_do_rebind(struct ctx *c, bool outbound)
+{
+	const uint8_t *fmap = outbound ? c->tcp.fwd_out.map : c->tcp.fwd_in.map;
+	const uint8_t *rmap = outbound ? c->tcp.fwd_in.map : c->tcp.fwd_out.map;
+	int (*socks)[IP_VERSIONS] = outbound ? tcp_sock_ns : tcp_sock_init_ext;
+	unsigned port;
+
+	for (port = 0; port < NUM_PORTS; port++) {
+		if (!bitmap_isset(fmap, port)) {
+			if (socks[port][V4] >= 0) {
+				close(socks[port][V4]);
+				socks[port][V4] = -1;
+			}
+
+			if (socks[port][V6] >= 0) {
+				close(socks[port][V6]);
+				socks[port][V6] = -1;
+			}
+
+			continue;
+		}
+
+		/* Don't loop back our own ports */
+		if (bitmap_isset(rmap, port))
+			continue;
+
+		if ((c->ifi4 && socks[port][V4] == -1) ||
+		    (c->ifi6 && socks[port][V6] == -1)) {
+			if (outbound)
+				tcp_ns_sock_init(c, port);
+			else
+				tcp_sock_init(c, AF_UNSPEC, NULL, NULL, port);
+		}
+	}
+}
+
 /**
  * struct tcp_port_rebind_arg - Arguments for tcp_port_rebind()
  * @c:			Execution context
@@ -3169,58 +3212,13 @@ struct tcp_port_rebind_arg {
 static int tcp_port_rebind(void *arg)
 {
 	struct tcp_port_rebind_arg *a = (struct tcp_port_rebind_arg *)arg;
-	unsigned port;
 
 	if (a->bind_in_ns) {
 		ns_enter(a->c);
 
-		for (port = 0; port < NUM_PORTS; port++) {
-			if (!bitmap_isset(a->c->tcp.fwd_out.map, port)) {
-				if (tcp_sock_ns[port][V4] >= 0) {
-					close(tcp_sock_ns[port][V4]);
-					tcp_sock_ns[port][V4] = -1;
-				}
-
-				if (tcp_sock_ns[port][V6] >= 0) {
-					close(tcp_sock_ns[port][V6]);
-					tcp_sock_ns[port][V6] = -1;
-				}
-
-				continue;
-			}
-
-			/* Don't loop back our own ports */
-			if (bitmap_isset(a->c->tcp.fwd_in.map, port))
-				continue;
-
-			if ((a->c->ifi4 && tcp_sock_ns[port][V4] == -1) ||
-			    (a->c->ifi6 && tcp_sock_ns[port][V6] == -1))
-				tcp_ns_sock_init(a->c, port);
-		}
+		tcp_port_do_rebind(a->c, true);
 	} else {
-		for (port = 0; port < NUM_PORTS; port++) {
-			if (!bitmap_isset(a->c->tcp.fwd_in.map, port)) {
-				if (tcp_sock_init_ext[port][V4] >= 0) {
-					close(tcp_sock_init_ext[port][V4]);
-					tcp_sock_init_ext[port][V4] = -1;
-				}
-
-				if (tcp_sock_init_ext[port][V6] >= 0) {
-					close(tcp_sock_init_ext[port][V6]);
-					tcp_sock_init_ext[port][V6] = -1;
-				}
-				continue;
-			}
-
-			/* Don't loop back our own ports */
-			if (bitmap_isset(a->c->tcp.fwd_out.map, port))
-				continue;
-
-			if ((a->c->ifi4 && tcp_sock_init_ext[port][V4] == -1) ||
-			    (a->c->ifi6 && tcp_sock_init_ext[port][V6] == -1))
-				tcp_sock_init(a->c, AF_UNSPEC, NULL, NULL,
-					      port);
-		}
+		tcp_port_do_rebind(a->c, false);
 	}
 
 	return 0;