Fix a bug where we could return too much data from the byte source. Add test.

diff --git a/bytesource.c b/bytesource.c
index 9e46d20ec0fdd8411326f097e38e774399075176..fed7551d50721c1e06ed2057c3a23e4dc83c069a 100644 (file)
--- a/bytesource.c
+++ b/bytesource.c
@@ -21,7 +21,7 @@ uint8_t byte_source_read_marker(struct byte_source* src)
 {
        // Refill until we have at least two bytes or EOF.
        while (src->bytes_available < 2) {
-               const unsigned bytes_to_read = BYTESOURCE_CHUNK_SIZE - src->bytes_available;
+               const unsigned bytes_to_read = 2 - src->bytes_available;
                const ssize_t bytes_read =
                        (*src->input_func)(src->userdata,
                                           src->bytes + src->bytes_available,
@@ -57,12 +57,13 @@ 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 <= (ssize_t)bytes_to_read);
                

diff --git a/bytesource_test.c b/bytesource_test.c
index 4dc7661c0a3e66cb150e6e95c1c6f935b39027dd..36c6da2926a2863c049d7fc5f4b6b41436da888d 100644 (file)
--- a/bytesource_test.c
+++ b/bytesource_test.c
@@ -148,6 +148,34 @@ void test_broken_marker()
        assert(ret == -1);
 }
 
+// Testing small reads -- even with a fast source, we shouldn't get more data
+// back than we asked for.
+void test_small_reads()
+{
+       uint8_t bytes[] = { 0xff, 0x80, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 };
+
+       struct custom_read_userdata ud;
+       ud.bytes = bytes;
+       ud.bytes_left = sizeof(bytes);
+
+       struct byte_source source;
+       init_byte_source(&source, custom_read, &ud);
+
+       uint8_t marker = byte_source_read_marker(&source);
+       assert(marker == 0x80);
+
+       uint8_t buf[4096];
+       for (unsigned i = 0; i < 8; ++i) {
+               ssize_t ret = byte_source_input_func(&source, buf, 1);
+               assert(ret == 1);
+               assert(buf[0] == i);
+       }
+       
+       // Now EOF.
+       ssize_t ret = byte_source_input_func(&source, buf, 4096);
+       assert(ret == 0);
+}
+
 int main(void)
 {
        init_choices();
@@ -161,6 +189,9 @@ int main(void)
        printf("test_broken_marker()\n");
        test_broken_marker();
        
+       printf("test_small_reads()\n");
+       test_small_reads();
+       
        printf("All tests pass.\n");
        return 0;
 }