From 6204e32f346d3c19ad5b95998b0838e29fcd2cd7 Mon Sep 17 00:00:00 2001
From: Miroslav Crnic <miroslav.crnic@xtxmarkets.com>
Date: Tue, 9 Dec 2025 16:25:15 +0000
Subject: [PATCH] kmod: refcount fix

---
 kmod/block.c | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/kmod/block.c b/kmod/block.c
index 40f8336a..efccc84a 100644
--- a/kmod/block.c
+++ b/kmod/block.c
@@ -686,13 +686,14 @@ static struct block_socket*  get_block_socket(
             ternfs_debug("multiple callers tried to get socket to %pI4:%d, dropping one", &other_sock->addr.sin_addr, ntohs(other_sock->addr.sin_port));
             // call again rather than trying to `sock_release` with the
             // RCU read lock held, this might not be safe in atomic context.
-            sock_release(sock->sock);
-            kfree(sock);
+            block_socket_put(sock);
             return get_block_socket(ops, addr);
         }
     }
     rcu_read_unlock();
 
+    block_socket_hold(sock); // we are now sure we'll return it
+
     // Put the new callbacks in
 
     sock->saved_data_ready = sock->sock->sk->sk_data_ready;
@@ -701,8 +702,6 @@ static struct block_socket*  get_block_socket(
     sock->sock->sk->sk_state_change = block_socket_sk_state_change;
     sock->sock->sk->sk_write_space = ops->sk_write_space;
 
-    block_socket_hold(sock); // we are now sure we'll return it
-
     // Insert the socket into the hash map -- anyone else which
     // will find it will be good to do.
     hlist_add_head_rcu(&sock->hnode, &ops->sockets[bucket]);
@@ -736,7 +735,6 @@ static int block_socket_receive_req_locked(
     size_t len0 = len;
 
     if (atomic_read(&socket->err)) {
-        block_socket_put(socket);
         return len0;
     }
 
@@ -754,6 +752,7 @@ static int block_socket_receive_req_locked(
         }
         int consumed = receive_single_req(req, skb, offset, len);
         if (consumed < 0) {
+            block_socket_hold(socket);
             error_socket(socket, consumed);
             return len0;
         }
@@ -786,7 +785,6 @@ static int block_socket_receive_req_locked(
 
     // We have more data but no requests. This should not happen
     BUG_ON(req == NULL && len > 0);
-    block_socket_put(socket);
     return len0 - len;
 }
 
@@ -803,7 +801,6 @@ static void block_socket_sk_data_ready(
         rd_desc.count = 1;
         // while this is protected by callback lock the callback might want to error a socket
         // and enqueue cleanup work which needs a reference. best to take it here
-        block_socket_hold(socket);
         tcp_read_sock(sk, &rd_desc, ops->receive_req);
         block_socket_state_check_locked(sk);
     }