ksmbd: update tree disconnect race update

Signed-off-by: Namjae Jeon <[email protected]>

mgmt/tree_connect.c

@@ -75,6 +75,7 @@ ksmbd_tree_conn_connect(struct ksmbd_conn *conn, struct ksmbd_session *sess,
 	tree_conn->share_conf = sc;
 	tree_conn->t_state = TREE_NEW;
 	status.tree_conn = tree_conn;
+	atomic_set(&tree_conn->refcount, 1);
 	init_waitqueue_head(&tree_conn->refcount_q);
 
 	ret = xa_err(xa_store(&sess->tree_conns, tree_conn->id, tree_conn,
@@ -117,7 +118,7 @@ int ksmbd_tree_conn_disconnect(struct ksmbd_session *sess,
 	xa_erase(&sess->tree_conns, tree_conn->id);
 	write_unlock(&sess->tree_conns_lock);
 
-	if (atomic_read(&tree_conn->refcount))
+	if (!atomic_dec_and_test(&tree_conn->refcount))
 		wait_event(tree_conn->refcount_q,
 			   atomic_read(&tree_conn->refcount) == 0);
 
@@ -138,8 +139,8 @@ struct ksmbd_tree_connect *ksmbd_tree_conn_lookup(struct ksmbd_session *sess,
 	if (tcon) {
 		if (tcon->t_state != TREE_CONNECTED)
 			tcon = NULL;
-		else
-			atomic_inc(&tcon->refcount);
+		else if (!atomic_inc_not_zero(&tcon->refcount))
+			tcon = NULL;
 	}
 	read_unlock(&sess->tree_conns_lock);