X-Git-Url: https://git.sesse.net/?p=fjl;a=blobdiff_plain;f=bytesource.c;h=c5e0ed88b1b0fd464e7ceb0849d0f78c98a11aad;hp=4ebe79c8d3e5afc34b3fed144ac28a5f6fef7bcb;hb=d5b9f8f45a79e6f25045bcc10511884a20fc8912;hpb=352c52d38c212a2aa93a9970731c9d7a2f8d4dac

diff --git a/bytesource.c b/bytesource.c
index 4ebe79c..c5e0ed8 100644
--- a/bytesource.c
+++ b/bytesource.c
@@ -1,4 +1,5 @@
 #include <stdbool.h>
+#include <stdlib.h>
 #include <string.h>
 #include <assert.h>
 
@@ -17,17 +18,34 @@ void init_byte_source(struct byte_source* source, raw_input_func_t* input_func,
 	source->userdata = userdata;
 }
 
-uint8_t byte_source_read_marker(struct byte_source* source)
+uint8_t byte_source_read_marker(struct byte_source* src)
 {
-	assert(source->bytes_available >= 2);
-	assert(source->bytes[0] == MARKER_CHAR);
-	assert(source->bytes[1] != STUFF_MARKER);
+	// Refill until we have at least two bytes or EOF.
+	while (src->bytes_available < 2) {
+		const unsigned bytes_to_read = 2 - src->bytes_available;
+		const ssize_t bytes_read =
+			(*src->input_func)(src->userdata,
+		                           src->bytes + src->bytes_available,
+			                   bytes_to_read);
+		assert(bytes_read >= -1);
+		assert(bytes_read <= (ssize_t)bytes_to_read);
+		
+		if (bytes_read == -1 || bytes_read == 0) {
+			return 0x00;
+		}
 
-	uint8_t ret = source->bytes[1];
+		src->bytes_available += bytes_read;
+	}
 
-	memmove(source->bytes, source->bytes + 2, source->bytes_available - 2);
-	source->bytes_available -= 2;
+	assert(src->bytes_available >= 2);
+	if (src->bytes[0] != MARKER_CHAR || src->bytes[1] == STUFF_MARKER) {
+		return 0x00;
+	}
 
+	uint8_t ret = src->bytes[1];
+	memmove(src->bytes, src->bytes + 2, src->bytes_available - 2);
+	src->bytes_available -= 2;
+	
 	return ret;
 }
 
@@ -40,14 +58,15 @@ ssize_t byte_source_input_func(void* source, uint8_t* buf, size_t len)
 	while (src->bytes_available == 0 ||
 	       (src->bytes_available == 1 && src->bytes[0] == MARKER_CHAR)) {
 		const unsigned space_left = BYTESOURCE_CHUNK_SIZE - src->bytes_available;
-		const size_t bytes_to_read = (len > space_left ? space_left : len);
+		const unsigned missing_data = len - src->bytes_available;
+		const size_t bytes_to_read = (missing_data > space_left ? space_left : missing_data);
 		assert(bytes_to_read <= BYTESOURCE_CHUNK_SIZE);
 		const ssize_t bytes_read =
 			(*src->input_func)(src->userdata,
-		                              src->bytes + src->bytes_available,
-			                      bytes_to_read);
+		                           src->bytes + src->bytes_available,
+			                   bytes_to_read);
 		assert(bytes_read >= -1);
-		assert(bytes_read <= bytes_to_read);
+		assert(bytes_read <= (ssize_t)bytes_to_read);
 		
 		if (bytes_read == -1) {
 			return -1;
@@ -63,7 +82,7 @@ ssize_t byte_source_input_func(void* source, uint8_t* buf, size_t len)
 
 		src->bytes_available += bytes_read;
 	}
-	
+
 	// Now unstuff as much as we can. First of all, if there's a 0xFF at the
 	// end of the buffer, we don't include it this time; the unstuff function
 	// will only give us an error since it can't decide if it's a marker or
@@ -76,45 +95,47 @@ ssize_t byte_source_input_func(void* source, uint8_t* buf, size_t len)
 		end_marker = true;
 	}
 
-	int unstuffed_bytes = (*unstuff_choice)(buf, src->bytes, bytes_to_unstuff);
-	assert(unstuffed_bytes != 0);
-	if (unstuffed_bytes > 0) {
-		// Fast path: No markers in the data. We can basically just
-		// return it.
-		if (end_marker) {
-			src->bytes_available = 1;
-			src->bytes[0] = 0xff;
-		} else {
-			src->bytes_available = 0;
-			src->bytes[0] = 0;
+	// Fast path: No markers in the data, we don't have more data than we
+	// need. We can basically just return it.
+	if (bytes_to_unstuff <= len) {
+		int unstuffed_bytes = (*unstuff_choice)(buf, src->bytes, bytes_to_unstuff);
+		assert(unstuffed_bytes != 0);
+		if (unstuffed_bytes > 0) {
+			// Fast path: No markers in the data. We can basically just
+			// return it.
+			if (end_marker) {
+				src->bytes_available = 1;
+				src->bytes[0] = 0xff;
+			} else {
+				src->bytes_available = 0;
+			}
+			return unstuffed_bytes;
 		}
-		return unstuffed_bytes;
 	}
 
 	// Slow path: There was a marker in the data. Unstuff manually until
 	// we hit the marker, then return that.
-	assert(unstuffed_bytes == -1);
 	unsigned bytes_read;
 	unsigned bytes_written = 0;
-	for (bytes_read = 0; bytes_read < src->bytes_available; ++bytes_read) {
-		buf[bytes_written++] = src->bytes[bytes_read];
+	for (bytes_read = 0; bytes_read < bytes_to_unstuff && bytes_written < len; ++bytes_read) {
 		if (src->bytes[bytes_read] != MARKER_CHAR) {
+			buf[bytes_written++] = src->bytes[bytes_read];
 			continue;
 		}
 
 		assert(bytes_read < src->bytes_available);
 		if (src->bytes[bytes_read + 1] == STUFF_MARKER) {
+			src->bytes[bytes_written++] = MARKER_CHAR;
+
 			// Skip the stuff byte.
 			++bytes_read;
-			continue;
 		} else {
 			// OK, this is our marker.
 			break;
-		}	
+		}
 	}
 
 	memmove(src->bytes, src->bytes + bytes_read, src->bytes_available - bytes_read);
 	src->bytes_available -= bytes_read;
-	assert(bytes_written >= 1);
-	return bytes_written - 1;
+	return bytes_written;
 }