Kernel: Hold socket back from accept() until it's fully set up

3 files changed, 21 insertions, 2 deletions

diff --git a/Kernel/Net/NetworkTask.cpp b/Kernel/Net/NetworkTask.cpp
index 199f1d67d8..85dd930ada 100644
--- a/Kernel/Net/NetworkTask.cpp
+++ b/Kernel/Net/NetworkTask.cpp
@@ -459,6 +459,7 @@ void handle_tcp(const IPv4Packet& ipv4_packet)
                 socket->set_setup_state(Socket::SetupState::Completed);
                 socket->set_connected(true);
             }
+            socket->release_to_originator();
             return;
         default:
             kprintf("handle_tcp: unexpected flags in SynReceived state\n");
diff --git a/Kernel/Net/TCPSocket.cpp b/Kernel/Net/TCPSocket.cpp
index 18ee8225fa..eb08ac1068 100644
--- a/Kernel/Net/TCPSocket.cpp
+++ b/Kernel/Net/TCPSocket.cpp
@@ -79,14 +79,27 @@ RefPtr<TCPSocket> TCPSocket::create_client(const IPv4Address& new_local_address,
     client->set_peer_address(new_peer_address);
     client->set_peer_port(new_peer_port);
     client->set_direction(Direction::Incoming);
+    client->set_originator(*this);
 
-    queue_connection_from(client);
-
+    m_pending_release_for_accept.set(tuple, client);
     sockets_by_tuple().resource().set(tuple, client);
 
     return from_tuple(tuple);
 }
 
+void TCPSocket::release_to_originator()
+{
+    ASSERT(!!m_originator);
+    m_originator->release_for_accept(this);
+}
+
+void TCPSocket::release_for_accept(RefPtr<TCPSocket> socket)
+{
+    ASSERT(m_pending_release_for_accept.contains(socket->tuple()));
+    m_pending_release_for_accept.remove(socket->tuple());
+    queue_connection_from(*socket);
+}
+
 TCPSocket::TCPSocket(int protocol)
     : IPv4Socket(SOCK_STREAM, protocol)
 {
diff --git a/Kernel/Net/TCPSocket.h b/Kernel/Net/TCPSocket.h
index f13591c3a2..adef5d9802 100644
--- a/Kernel/Net/TCPSocket.h
+++ b/Kernel/Net/TCPSocket.h
@@ -126,6 +126,9 @@ public:
     static RefPtr<TCPSocket> from_endpoints(const IPv4Address& local_address, u16 local_port, const IPv4Address& peer_address, u16 peer_port);
 
     RefPtr<TCPSocket> create_client(const IPv4Address& local_address, u16 local_port, const IPv4Address& peer_address, u16 peer_port);
+    void set_originator(RefPtr<TCPSocket> originator) { m_originator = originator; }
+    void release_to_originator();
+    void release_for_accept(RefPtr<TCPSocket>);
 
 protected:
     void set_direction(Direction direction) { m_direction = direction; }
@@ -144,6 +147,8 @@ private:
     virtual KResult protocol_bind() override;
     virtual KResult protocol_listen() override;
 
+    RefPtr<TCPSocket> m_originator;
+    HashMap<IPv4SocketTuple, NonnullRefPtr<TCPSocket>> m_pending_release_for_accept;
     Direction m_direction { Direction::Unspecified };
     Error m_error { Error::None };
     WeakPtr<NetworkAdapter> m_adapter;