From 35e1df6bb4e2a438714f1d0a9f1116df94927859 Mon Sep 17 00:00:00 2001
From: Jukka Rissanen <jukka.rissanen@nordicsemi.no>
Date: Tue, 2 Jan 2024 17:46:30 +0200
Subject: [PATCH] tests: net: tcp: Add support for close callback

Add a function callback that is called when the TCP connection
is closed. This is only available if doing network tests.

Signed-off-by: Jukka Rissanen <jukka.rissanen@nordicsemi.no>
---
 subsys/net/ip/tcp.c         | 18 ++++++++++++++++++
 subsys/net/ip/tcp_private.h | 13 +++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/subsys/net/ip/tcp.c b/subsys/net/ip/tcp.c
index c4db136f37..41df924607 100644
--- a/subsys/net/ip/tcp.c
+++ b/subsys/net/ip/tcp.c
@@ -714,6 +714,12 @@ static void tcp_conn_release(struct k_work *work)
 	struct tcp *conn = CONTAINER_OF(work, struct tcp, conn_release);
 	struct net_pkt *pkt;
 
+#if defined(CONFIG_NET_TEST)
+	if (conn->test_closed_cb != NULL) {
+		conn->test_closed_cb(conn, conn->test_user_data);
+	}
+#endif
+
 	k_mutex_lock(&tcp_lock, K_FOREVER);
 
 	/* Application is no longer there, unref any remaining packets on the
@@ -762,6 +768,18 @@ static void tcp_conn_release(struct k_work *work)
 	k_mutex_unlock(&tcp_lock);
 }
 
+#if defined(CONFIG_NET_TEST)
+void tcp_install_close_cb(struct net_context *ctx,
+			  net_tcp_closed_cb_t cb,
+			  void *user_data)
+{
+	NET_ASSERT(ctx->tcp != NULL);
+
+	((struct tcp *)ctx->tcp)->test_closed_cb = cb;
+	((struct tcp *)ctx->tcp)->test_user_data = user_data;
+}
+#endif
+
 static int tcp_conn_unref(struct tcp *conn)
 {
 	int ref_count = atomic_get(&conn->ref_count);
diff --git a/subsys/net/ip/tcp_private.h b/subsys/net/ip/tcp_private.h
index 5d5d78d916..1f2eee2498 100644
--- a/subsys/net/ip/tcp_private.h
+++ b/subsys/net/ip/tcp_private.h
@@ -250,6 +250,9 @@ struct tcp_collision_avoidance_reno {
 };
 #endif
 
+struct tcp;
+typedef void (*net_tcp_closed_cb_t)(struct tcp *conn, void *user_data);
+
 struct tcp { /* TCP connection */
 	sys_snode_t next;
 	struct net_context *context;
@@ -263,6 +266,10 @@ struct tcp { /* TCP connection */
 		struct tcp *accepted_conn;
 	};
 	net_context_connect_cb_t connect_cb;
+#if defined(CONFIG_NET_TEST)
+	net_tcp_closed_cb_t test_closed_cb;
+	void *test_user_data;
+#endif
 	struct k_mutex lock;
 	struct k_sem connect_sem; /* semaphore for blocking connect */
 	struct k_sem tx_sem; /* Semaphore indicating if transfers are blocked . */
@@ -345,3 +352,9 @@ struct tcp { /* TCP connection */
 	_flags(_fl, _op, _mask, sizeof(#_args) > 1 ? _args : true)
 
 typedef void (*net_tcp_cb_t)(struct tcp *conn, void *user_data);
+
+#if defined(CONFIG_NET_TEST)
+void tcp_install_close_cb(struct net_context *ctx,
+			  net_tcp_closed_cb_t cb,
+			  void *user_data);
+#endif