Start by adding a packet size to the existing protocol

1 files changed, 54 insertions, 9 deletions

diff --git a/sockutil.c b/sockutil.c
index 14a058d..4ccca03 100644
--- a/sockutil.c
+++ b/sockutil.c
@@ -19,6 +19,7 @@
  */
 #include <sys/types.h>
 #include <sys/socket.h>
+#include <sys/uio.h>
 #include <netinet/in.h>
 #include <stdio.h>
 #include <string.h>
@@ -196,22 +197,58 @@ connect_to_peer(struct sockaddr_storage *peer, int port)
 	return sd;
 }
 
+/* 
+ * All incoming and outgoing packets go through this function.
+ * Caller should make sure the fd is ready to read.
+ */
+typedef struct {
+	char buf[QARSH_MAX_PACKET_SIZE];
+	int start;
+	int end;
+} Buffer;
 
 struct qa_packet *
 recv_packet(int fd)
 {
-	char *packetbuf;
-	int packetsize;
+	static Buffer pb = { "", 0, 0 };
+	uint32_t packetsize;
 	struct qa_packet *qp = NULL;
+	int ret;
 
-	packetbuf = malloc(QARSH_MAX_PACKET_SIZE);
-	memset(packetbuf, 0, QARSH_MAX_PACKET_SIZE);
+	if (pb.start == pb.end && pb.start)
+		pb.start = pb.end = 0;
+
+recv_read:
+	ret = read(fd, pb.buf, QARSH_MAX_PACKET_SIZE - pb.start);
+	if (ret > 0) {
+		pb.end = pb.start + ret;
+		/* get packet size from the packet */
+		packetsize = ntohl(*(uint32_t *)(pb.buf+pb.start));
+		if (packetsize > QARSH_MAX_PACKET_SIZE) {
+			printf("Packet size too large, %d > %d\n", packetsize, QARSH_MAX_PACKET_SIZE);
+			return NULL;
+		} else if (packetsize <= (ret - sizeof packetsize)) {
+			pb.start += sizeof packetsize;
+			qp = parse_packets(pb.buf + pb.start, pb.end - pb.start);
+		} else { /* packetsize > what we read, we need to read more */
+			printf("Compacting buffer to fit rest of packet, %d read, %d left\n", ret, packetsize - ret);
+			memmove(pb.buf, pb.buf + pb.start, pb.end - pb.start);
+			pb.end -= pb.start;
+			pb.start = 0;
+			goto recv_read;
+		}
 
-	packetsize = read(fd, packetbuf, QARSH_MAX_PACKET_SIZE);
-	if (packetsize > 0) {
-		qp = parse_packets(packetbuf, packetsize);
+		/* Get the packet buffer ready for the next call */
+		if (packetsize == (ret - sizeof packetsize)) { /* The read was a complete packet */
+			pb.start = pb.end = 0;
+			printf("Complete read used\n");
+		} else {
+			pb.start += packetsize;
+			memmove(pb.buf, pb.buf + pb.start, pb.end - pb.start);
+			pb.end -= pb.start;
+			pb.start = 0;
+		}
 	}
-	free(packetbuf);
 	return qp;
 }
 
@@ -220,11 +257,19 @@ send_packet(int fd, struct qa_packet *qp)
 {
 	char *packetbuf;
 	int packetsize;
+	uint32_t netsize;
 	ssize_t ret;
+	struct iovec iovs[2];
 
 	packetbuf = qptostr(qp, &packetbuf, &packetsize);
 
-	ret = write(fd, packetbuf, packetsize);
+	netsize = htonl(packetsize);
+	iovs[0].iov_base = &netsize;
+	iovs[0].iov_len = sizeof netsize;
+	iovs[1].iov_base = packetbuf;
+	iovs[1].iov_len = packetsize;
+
+	ret = writev(fd, iovs, 2);
 	free(packetbuf);
 	return ret;
 }