[PATCH 7/8] trace-cmd: Add tracecmd_msg_wait_close function

[Slavomir Kaslev <kaslevs@xxxxxxxxxx>]

Add tracecmd_msg_wait_close function that waits for MSG_CLOSE and logs any
invalid messages it receives.

Also switch tracecmd_msg_collect_data to use the new function while at it.

Signed-off-by: Slavomir Kaslev <kaslevs@xxxxxxxxxx>
---
 include/trace-cmd/trace-cmd.h |  1 +
 tracecmd/trace-msg.c          | 35 +++++++++++++++--------------------
 2 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/include/trace-cmd/trace-cmd.h b/include/trace-cmd/trace-cmd.h
index 0ab23f6..ca4452b 100644
--- a/include/trace-cmd/trace-cmd.h
+++ b/include/trace-cmd/trace-cmd.h
@@ -319,6 +319,7 @@ int tracecmd_msg_data_send(struct tracecmd_msg_handle *msg_handle,
 			       const char *buf, int size);
 int tracecmd_msg_finish_sending_data(struct tracecmd_msg_handle *msg_handle);
 int tracecmd_msg_send_close_msg(struct tracecmd_msg_handle *msg_handle);
+int tracecmd_msg_wait_close(struct tracecmd_msg_handle *msg_handle);
 
 /* for server */
 int tracecmd_msg_initial_setting(struct tracecmd_msg_handle *msg_handle);
diff --git a/tracecmd/trace-msg.c b/tracecmd/trace-msg.c
index c24424b..5079d43 100644
--- a/tracecmd/trace-msg.c
+++ b/tracecmd/trace-msg.c
@@ -659,39 +659,34 @@ error:
 
 int tracecmd_msg_collect_data(struct tracecmd_msg_handle *msg_handle, int ofd)
 {
-	struct tracecmd_msg msg;
-	u32 cmd;
 	int ret;
 
 	ret = tracecmd_msg_read_data(msg_handle, ofd);
 	if (ret)
-		goto error;
+		return ret;
+
+	return tracecmd_msg_wait_close(msg_handle);
+}
 
-	/* check the finish message of the client */
+int tracecmd_msg_wait_close(struct tracecmd_msg_handle *msg_handle)
+{
+	struct tracecmd_msg msg;
+	int ret = -1;
+
+	memset(&msg, 0, sizeof(msg));
 	while (!tracecmd_msg_done(msg_handle)) {
 		ret = tracecmd_msg_recv(msg_handle->fd, &msg);
-		if (ret < 0) {
-			warning("reading client");
-			return ret;
-		}
-
-		cmd = ntohl(msg.hdr.cmd);
-		if (cmd == MSG_CLOSE)
-			/* Finish this connection */
-			break;
-		else {
-			warning("Not accept the message %d", ntohl(msg.hdr.cmd));
-			ret = -EINVAL;
+		if (ret < 0)
 			goto error;
-		}
 
+		if (ntohl(msg.hdr.cmd) == MSG_CLOSE)
+			return 0;
+
+		error_operation(&msg);
 		msg_free(&msg);
 	}
 
-	return 0;
-
 error:
-	error_operation(&msg);
 	msg_free(&msg);
 	return ret;
 }
-- 
2.19.1