3 files changed, 109 insertions, 37 deletions

diff --git a/src/msg.c b/src/msg.c
index df8e14e..5e18681 100644
--- a/src/msg.c
+++ b/src/msg.c
@@ -102,7 +102,7 @@ int msg_from_argv(struct msg *msg, char **argv)
 	return msg_copy_argv(msg, argv);
 }
 
-static uint32_t calc_buf_len(const struct msg *msg)
+static uint32_t calc_buf_size(const struct msg *msg)
 {
 	uint32_t len = 0;
 	for (int i = 0; i < msg->argc; ++i)
@@ -157,22 +157,30 @@ free:
 
 int msg_send(int fd, const struct msg *msg)
 {
+	struct buf *buf;
 	int ret = 0;
 
-	uint32_t len = calc_buf_len(msg);
-	char *buf = malloc(len);
-	if (!buf) {
+	uint32_t size = calc_buf_size(msg);
+	char *data = malloc(size);
+	if (!data) {
 		log_errno("malloc");
 		return -1;
 	}
-	argv_pack(buf, msg);
+	argv_pack(data, msg);
 
-	ret = net_send_buf(fd, buf, len);
+	ret = buf_create(&buf, data, size);
 	if (ret < 0)
-		goto free_buf;
+		goto free_data;
 
-free_buf:
-	free(buf);
+	ret = net_send_buf(fd, buf);
+	if (ret < 0)
+		goto destroy_buf;
+
+destroy_buf:
+	buf_destroy(buf);
+
+free_data:
+	free(data);
 
 	return ret;
 }
@@ -194,22 +202,21 @@ int msg_send_and_wait(int fd, const struct msg *request, struct msg *response)
 
 int msg_recv(int fd, struct msg *msg)
 {
-	void *buf;
-	uint32_t len;
+	struct buf *buf;
 	int ret = 0;
 
-	ret = net_recv_buf(fd, &buf, &len);
+	ret = net_recv_buf(fd, &buf);
 	if (ret < 0)
 		return ret;
 
-	msg->argc = calc_argv_len(buf, len);
+	msg->argc = calc_argv_len(buf_get_data(buf), buf_get_size(buf));
 
-	ret = argv_unpack(msg, buf);
+	ret = argv_unpack(msg, buf_get_data(buf));
 	if (ret < 0)
-		goto free_buf;
+		goto destroy_buf;
 
-free_buf:
-	free(buf);
+destroy_buf:
+	buf_destroy(buf);
 
 	return ret;
 }
diff --git a/src/net.c b/src/net.c
index c972549..a1aa941 100644
--- a/src/net.c
+++ b/src/net.c
@@ -198,51 +198,108 @@ int net_recv_all(int fd, void *buf, size_t len)
 	return 0;
 }
 
-int net_send_buf(int fd, const void *buf, uint32_t len)
+struct buf {
+	uint32_t size;
+	void *data;
+};
+
+int buf_create(struct buf **_buf, const void *data, uint32_t size)
 {
+	struct buf *buf;
 	int ret = 0;
 
-	len = htonl(len);
-	ret = net_send_all(fd, &len, sizeof(len));
+	*_buf = malloc(sizeof(struct buf));
+	if (!*_buf) {
+		log_errno("malloc");
+		return -1;
+	}
+	buf = *_buf;
+
+	buf->data = malloc(size);
+	if (!buf->data) {
+		log_errno("malloc");
+		goto free;
+	}
+
+	buf->size = size;
+	memcpy(buf->data, data, size);
+
+	return ret;
+
+free:
+	free(buf);
+
+	return ret;
+}
+
+void buf_destroy(struct buf *buf)
+{
+	free(buf->data);
+	free(buf);
+}
+
+uint32_t buf_get_size(const struct buf *buf)
+{
+	return buf->size;
+}
+
+void *buf_get_data(const struct buf *buf)
+{
+	return buf->data;
+}
+
+int net_send_buf(int fd, const struct buf *buf)
+{
+	uint32_t size;
+	int ret = 0;
+
+	size = htonl(buf->size);
+	ret = net_send_all(fd, &size, sizeof(size));
 	if (ret < 0)
 		return ret;
-	len = ntohl(len);
 
-	ret = net_send_all(fd, buf, len);
+	ret = net_send_all(fd, buf->data, buf->size);
 	if (ret < 0)
 		return ret;
 
 	return ret;
 }
 
-int net_recv_buf(int fd, void **buf, uint32_t *len)
+int net_recv_buf(int fd, struct buf **buf)
 {
+	void *data;
+	uint32_t size;
 	int ret = 0;
 
-	ret = net_recv_all(fd, len, sizeof(*len));
+	ret = net_recv_all(fd, &size, sizeof(size));
 	if (ret < 0) {
-		log_err("Couldn't read buffer length\n");
+		log_err("Couldn't read buffer size\n");
 		goto fail;
 	}
+	size = ntohl(size);
 
-	*len = ntohl(*len);
-
-	*buf = malloc(*len);
-	if (!*buf) {
+	data = malloc(size);
+	if (!data) {
 		log_errno("malloc");
 		goto fail;
 	}
 
-	ret = net_recv_all(fd, *buf, *len);
+	ret = net_recv_all(fd, data, size);
 	if (ret < 0) {
 		log_err("Couldn't read buffer\n");
-		goto free_buf;
+		goto free_data;
 	}
 
-	return 0;
+	ret = buf_create(buf, data, size);
+	if (ret < 0)
+		goto free_data;
+
+	free(data);
+
+	return ret;
 
-free_buf:
-	free(*buf);
+free_data:
+	free(data);
 
 fail:
 	return -1;
diff --git a/src/net.h b/src/net.h
index b6dd174..185f10b 100644
--- a/src/net.h
+++ b/src/net.h
@@ -16,9 +16,17 @@ int net_accept(int fd);
 int net_connect(const char *host, const char *port);
 
 int net_send_all(int fd, const void *, size_t);
-int net_send_buf(int fd, const void *, uint32_t);
-
 int net_recv_all(int fd, void *, size_t);
-int net_recv_buf(int fd, void **, uint32_t *);
+
+struct buf;
+
+int buf_create(struct buf **, const void *, uint32_t);
+void buf_destroy(struct buf *);
+
+uint32_t buf_get_size(const struct buf *);
+void *buf_get_data(const struct buf *);
+
+int net_send_buf(int fd, const struct buf *);
+int net_recv_buf(int fd, struct buf **);
 
 #endif