make puzbus slightly more robust

3 files changed, 62 insertions, 21 deletions

diff --git a/client/main.cpp b/client/main.cpp
index 30d7045..dcc965b 100644
--- a/client/main.cpp
+++ b/client/main.cpp
@@ -17,31 +17,44 @@ int send_message() {
 	size_t size;
 	if (!pb_write(&output, &packed, &size)) {
 		printf("error writing!\n");
-		return 1;
+		return EXIT_FAILURE;
 	}
 
 	fwrite(packed, sizeof(packed[0]), size, stdout);
 	fflush(stdout);
 
-	return 0;
+	return EXIT_SUCCESS;
 }
 
 int read_message() {
 	freopen(NULL, "rb", stdin); // allow binary on stdin
 	struct pb_msg input;
 	
-	char buf[8]; // extremely small buffer to test chunked message parsing
+	char buf[4]; // extremely small buffer to test chunked message parsing
 	size_t bytes = 0;
+
 	while ((bytes = fread(buf, sizeof(buf[0]), sizeof(buf), stdin)) > 0) {
-		if (!pb_read(&input, buf, bytes)) continue;
+		int ret = pb_read(&input, buf, bytes);
+
+		// header read error
+		if (ret < 0) {
+			printf("error reading!\n");
+			return EXIT_FAILURE;
+		}
+
+		// continue reading if more bytes needed...
+		if (ret > 0) continue;
 
+		// message read completely!
 		printf("address: 0x%02x\n", input.addr);
 		printf("data:    \"%.*s\"\n", input.length, input.data);
 		free(input.data);
-		return 0;
+		return EXIT_SUCCESS;
 	}
 
-	return 1;
+	// if we reach this point, data was read but it did not contain a complete
+	// message, and is thus considered a failure
+	return EXIT_FAILURE;
 }
 
 int main() {
@@ -49,6 +62,6 @@ int main() {
 	if (!isatty(fileno(stdin))) return read_message();
 	
 	printf("please pipe some data in or out to use this program\n");
-	return 0;
+	return EXIT_SUCCESS;
 }
 
diff --git a/proto/puzbusv1.c b/proto/puzbusv1.c
index d6cd597..3be4939 100644
--- a/proto/puzbusv1.c
+++ b/proto/puzbusv1.c
@@ -7,21 +7,32 @@
 
 #include "puzbusv1.h"
 
-bool pb_read(struct pb_msg* target, char* buf, size_t buf_sz) {
-	// remaining bytes to be read to target->data; this is the only variable that
-	// needs to persist between buffer blocks, and is therefore static
-	static size_t rdata = 0;
+/**
+ * \brief Remaining bytes to be read to target->data
+ *
+ * This is the only variable that needs to persist between buffer blocks. It is
+ * declared in the global scope to allow resetting using the \c pb_read_reset()
+ * function.
+ *
+ * \note \p rdata may be reset by calling \c pb_read_reset()
+ */
+static size_t rdata = 0;
 
+int pb_read(struct pb_msg* target, char* buf, size_t buf_sz) {
 	// a new reader is used per buffer block passed to this function
 	mpack_reader_t reader;
 	mpack_reader_init_data(&reader, buf, buf_sz);
 
 	// at start of message
 	if (rdata == 0) {
-		// NOTE: This approach will crash and burn when target->addr can be read
-		// and target->length is past the end of the current buffer block. This is
-		// a highly unlikely scenario, as pb_read is called for each chunk of a TCP
-		// frame, and frames (should) include only one puzzle bus message.
+		// NOTE: The entire start of a message needs to be readable from the buffer
+		// at this point. When target->addr can be read and target->length is past
+		// the end of the current buffer block, this function will crash and burn.
+		// This is a highly unlikely scenario, as pb_read is called for each chunk
+		// of a TCP frame, and frames (should) include only one puzzle bus message.
+		// The check here is kind of optional.
+		if (buf_sz < 4) return -1;
+
 		target->addr = mpack_expect_u16(&reader);
 		target->length = rdata = mpack_expect_bin(&reader);
 		target->data = (char*) malloc(target->length);
@@ -35,7 +46,11 @@ bool pb_read(struct pb_msg* target, char* buf, size_t buf_sz) {
 	rdata -= to_read;
 
 	// if rdata = 0, the message was completely read
-	return rdata == 0;
+	return rdata;
+}
+
+void pb_read_reset() {
+	rdata = 0;
 }
 
 bool pb_write(struct pb_msg* target, char** buf, size_t* buf_sz) {
diff --git a/proto/puzbusv1.h b/proto/puzbusv1.h
index 814bc93..95130c9 100644
--- a/proto/puzbusv1.h
+++ b/proto/puzbusv1.h
@@ -24,13 +24,26 @@ struct pb_msg {
  * \param buf     pointer to input stream data chunk
  * \param buf_sz  size of \p buf
  *
- * \returns boolean true if a message was completely read, false if more data
- * is required
+ * \returns Integer representing amount of bytes required to finish message, or
+ * -1 if the message header could not be read. If this function returns 0, the
+ *  message in \p target is complete.
  *
- * \note target->data will automatically be allocated by this function, and
- * must be `free()`d by the caller when finished
+ * \note target->data will automatically be allocated by this function, even if
+ * the message is not fully parsed. This variable must be `free()`d by the
+ * caller after each complete message to prevent memory leaks.
  */
-bool pb_read(struct pb_msg* target, char* buf, size_t buf_sz);
+int pb_read(struct pb_msg* target, char* buf, size_t buf_sz);
+
+/**
+ * \brief reset the remaining message data counter
+ *
+ * Calling this function has the effect of forcing \c pb_read() to parse the
+ * next buffer chunk as the start of a new message. This function may be called
+ * before reading a TCP frame's data to mitigate any synchronization issues
+ * arising from earlier corrupt or otherwise malformed messages.
+ */
+void pb_read_reset();
+
 /**
  * \brief Allocate and write a msgpack-formatted message to \p buf
  *