Kernel: Lock TCPSocket lookup table across destruction

Use the same trick as SlavePTY and override unref() to provide safe
removal from the sockets_by_tuple table when destroying a TCPSocket.

This should fix the TCPSocket::from_tuple() flake seen on CI.

2 files changed, 17 insertions, 0 deletions

diff --git a/Kernel/Net/TCPSocket.cpp b/Kernel/Net/TCPSocket.cpp
index fa7eb8cc10..cf4c6f71ff 100644
--- a/Kernel/Net/TCPSocket.cpp
+++ b/Kernel/Net/TCPSocket.cpp
@@ -29,6 +29,21 @@ void TCPSocket::for_each(Function<void(const TCPSocket&)> callback)
     });
 }
 
+bool TCPSocket::unref() const
+{
+    bool did_hit_zero = sockets_by_tuple().with_exclusive([&](auto& table) {
+        if (deref_base())
+            return false;
+        table.remove(tuple());
+        return true;
+    });
+    if (did_hit_zero) {
+        const_cast<TCPSocket&>(*this).will_be_destroyed();
+        delete this;
+    }
+    return did_hit_zero;
+}
+
 void TCPSocket::set_state(State new_state)
 {
     dbgln_if(TCP_SOCKET_DEBUG, "TCPSocket({}) state moving from {} to {}", this, to_string(m_state), to_string(new_state));
diff --git a/Kernel/Net/TCPSocket.h b/Kernel/Net/TCPSocket.h
index 4cf0f32296..85f4c9ec98 100644
--- a/Kernel/Net/TCPSocket.h
+++ b/Kernel/Net/TCPSocket.h
@@ -22,6 +22,8 @@ public:
     static ErrorOr<NonnullRefPtr<TCPSocket>> try_create(int protocol, NonnullOwnPtr<DoubleBuffer> receive_buffer);
     virtual ~TCPSocket() override;
 
+    virtual bool unref() const override;
+
     enum class Direction {
         Unspecified,
         Outgoing,