udebug: wait for response after buffer add/remove

Fixes a race condition where freeing a buffer and immediately re-allocating and
adding it would fail to pass the file descriptor to udebugd

Signed-off-by: Felix Fietkau <nbd@nbd.name>

3 files changed, 88 insertions, 79 deletions

diff --git a/udebug-priv.h b/udebug-priv.h
index 9e43ddf..96373f9 100644
--- a/udebug-priv.h
+++ b/udebug-priv.h
@@ -25,9 +25,8 @@
 #define UDEBUG_TIMEOUT	1000
 
 __hidden int udebug_buf_open(struct udebug_buf *buf, int fd, uint32_t ring_size, uint32_t data_size);
-__hidden struct udebug_client_msg *__udebug_poll(struct udebug *ctx, int *fd, bool wait);
-__hidden void udebug_send_msg(struct udebug *ctx, struct udebug_client_msg *msg,
-		     struct blob_attr *meta, int fd);
+__hidden struct udebug_client_msg *
+udebug_send_and_wait(struct udebug *ctx, struct udebug_client_msg *msg, int *rfd);
 
 static inline int32_t u32_sub(uint32_t a, uint32_t b)
 {
diff --git a/udebug-remote.c b/udebug-remote.c
index f13f633..0b797f9 100644
--- a/udebug-remote.c
+++ b/udebug-remote.c
@@ -17,31 +17,6 @@
  */
 #include "udebug-priv.h"
 
-static struct udebug_client_msg *
-send_and_wait(struct udebug *ctx, struct udebug_client_msg *msg, int *rfd)
-{
-	int type = msg->type;
-	int fd = -1;
-
-	udebug_send_msg(ctx, msg, NULL, -1);
-
-	do {
-		if (fd >= 0)
-			close(fd);
-		fd = -1;
-		msg = __udebug_poll(ctx, &fd, true);
-	} while (msg && msg->type != type);
-	if (!msg)
-		return NULL;
-
-	if (rfd)
-		*rfd = fd;
-	else if (fd >= 0)
-		close(fd);
-
-	return msg;
-}
-
 static int
 udebug_remote_get_handle(struct udebug *ctx)
 {
@@ -53,7 +28,7 @@ udebug_remote_get_handle(struct udebug *ctx)
 	if (ctx->poll_handle >= 0 || !udebug_is_connected(ctx))
 		return 0;
 
-	msg = send_and_wait(ctx, &send_msg, NULL);
+	msg = udebug_send_and_wait(ctx, &send_msg, NULL);
 	if (!msg)
 		return -1;
 
@@ -82,7 +57,7 @@ int udebug_remote_buf_map(struct udebug *ctx, struct udebug_remote_buf *rb, uint
 	if (rb->buf.data || !udebug_is_connected(ctx))
 		return -1;
 
-	msg = send_and_wait(ctx, &send_msg, &fd);
+	msg = udebug_send_and_wait(ctx, &send_msg, &fd);
 	if (!msg || fd < 0)
 		return -1;
 
diff --git a/udebug.c b/udebug.c
index 00a66ab..16beefd 100644
--- a/udebug.c
+++ b/udebug.c
@@ -290,8 +290,9 @@ recv_retry(int fd, struct iovec *iov, bool wait, int *recv_fd)
 	return total;
 }
 
-void udebug_send_msg(struct udebug *ctx, struct udebug_client_msg *msg,
-		     struct blob_attr *meta, int fd)
+static void
+udebug_send_msg(struct udebug *ctx, struct udebug_client_msg *msg,
+		struct blob_attr *meta, int fd)
 {
 	struct iovec iov[2] = {
 		{ .iov_base = msg, .iov_len = sizeof(*msg) },
@@ -308,6 +309,79 @@ void udebug_send_msg(struct udebug *ctx, struct udebug_client_msg *msg,
 	writev_retry(ctx->fd.fd, iov, ARRAY_SIZE(iov), fd);
 }
 
+static bool
+udebug_recv_msg(struct udebug *ctx, struct udebug_client_msg *msg, int *fd,
+		bool wait)
+{
+	struct iovec iov = {
+		.iov_base = msg,
+		.iov_len = sizeof(*msg)
+	};
+	int ret;
+
+	ret = recv_retry(ctx->fd.fd, &iov, wait, fd);
+	if (ret == -2)
+		__udebug_disconnect(ctx, true);
+
+	return ret == sizeof(*msg);
+}
+
+static struct udebug_client_msg *
+__udebug_poll(struct udebug *ctx, int *fd, bool wait)
+{
+	static struct udebug_client_msg msg = {};
+
+	while (udebug_recv_msg(ctx, &msg, fd, wait)) {
+		struct udebug_remote_buf *rb;
+		void *key;
+
+		if (msg.type != CL_MSG_RING_NOTIFY)
+			return &msg;
+
+		if (fd && *fd >= 0)
+			close(*fd);
+
+		if (!ctx->notify_cb)
+			continue;
+
+		key = (void *)(uintptr_t)msg.id;
+		rb = avl_find_element(&ctx->remote_rings, key, rb, node);
+		if (!rb || !rb->poll)
+			continue;
+
+		if (ctx->poll_handle >= 0)
+			__atomic_fetch_or(&rb->buf.hdr->notify,
+					  1UL << ctx->poll_handle,
+					  __ATOMIC_RELAXED);
+		ctx->notify_cb(ctx, rb);
+	}
+
+	return NULL;
+}
+
+static struct udebug_client_msg *
+udebug_wait_for_response(struct udebug *ctx, struct udebug_client_msg *msg, int *rfd)
+{
+	int type = msg->type;
+	int fd = -1;
+
+	do {
+		if (fd >= 0)
+			close(fd);
+		fd = -1;
+		msg = __udebug_poll(ctx, &fd, true);
+	} while (msg && msg->type != type);
+	if (!msg)
+		return NULL;
+
+	if (rfd)
+		*rfd = fd;
+	else if (fd >= 0)
+		close(fd);
+
+	return msg;
+}
+
 static void
 udebug_buf_msg(struct udebug_buf *buf, enum udebug_client_msg_type type)
 {
@@ -317,6 +391,7 @@ udebug_buf_msg(struct udebug_buf *buf, enum udebug_client_msg_type type)
 	};
 
 	udebug_send_msg(buf->ctx, &msg, NULL, -1);
+	udebug_wait_for_response(buf->ctx, &msg, NULL);
 }
 
 static size_t __udebug_headsize(unsigned int ring_size, unsigned int page_size)
@@ -604,6 +679,7 @@ __udebug_buf_add(struct udebug *ctx, struct udebug_buf *buf)
 	blobmsg_close_array(&b, c);
 
 	udebug_send_msg(ctx, &msg, b.head, buf->fd);
+	udebug_wait_for_response(ctx, &msg, NULL);
 }
 
 int udebug_buf_add(struct udebug *ctx, struct udebug_buf *buf,
@@ -683,58 +759,17 @@ int udebug_connect(struct udebug *ctx, const char *path)
 	return 0;
 }
 
-static bool
-udebug_recv_msg(struct udebug *ctx, struct udebug_client_msg *msg, int *fd,
-		bool wait)
+void udebug_poll(struct udebug *ctx)
 {
-	struct iovec iov = {
-		.iov_base = msg,
-		.iov_len = sizeof(*msg)
-	};
-	int ret;
-
-	ret = recv_retry(ctx->fd.fd, &iov, wait, fd);
-	if (ret == -2)
-		__udebug_disconnect(ctx, true);
-
-	return ret == sizeof(*msg);
+	while (__udebug_poll(ctx, NULL, false));
 }
 
-struct udebug_client_msg *__udebug_poll(struct udebug *ctx, int *fd, bool wait)
+struct udebug_client_msg *
+udebug_send_and_wait(struct udebug *ctx, struct udebug_client_msg *msg, int *rfd)
 {
-	static struct udebug_client_msg msg = {};
-
-	while (udebug_recv_msg(ctx, &msg, fd, wait)) {
-		struct udebug_remote_buf *rb;
-		void *key;
-
-		if (msg.type != CL_MSG_RING_NOTIFY)
-			return &msg;
-
-		if (fd && *fd >= 0)
-			close(*fd);
-
-		if (!ctx->notify_cb)
-			continue;
-
-		key = (void *)(uintptr_t)msg.id;
-		rb = avl_find_element(&ctx->remote_rings, key, rb, node);
-		if (!rb || !rb->poll)
-			continue;
+	udebug_send_msg(ctx, msg, NULL, -1);
 
-		if (ctx->poll_handle >= 0)
-			__atomic_fetch_or(&rb->buf.hdr->notify,
-					  1UL << ctx->poll_handle,
-					  __ATOMIC_RELAXED);
-		ctx->notify_cb(ctx, rb);
-	}
-
-	return NULL;
-}
-
-void udebug_poll(struct udebug *ctx)
-{
-	while (__udebug_poll(ctx, NULL, false));
+	return udebug_wait_for_response(ctx, msg, rfd);
 }
 
 static void udebug_fd_cb(struct uloop_fd *fd, unsigned int events)