X-Git-Url: https://git.sesse.net/?p=fjl;a=blobdiff_plain;f=bitsource.c;h=b1593aa6ec554959b4bba6bbe3aceb75d4e04863;hp=6cc627c70b08b462b0d0c9c7e0a4ec3d25cdceed;hb=HEAD;hpb=ddc3c24c837f081787825d128675eca628d0e91c

diff --git a/bitsource.c b/bitsource.c
index 6cc627c..b1593aa 100644
--- a/bitsource.c
+++ b/bitsource.c
@@ -8,12 +8,14 @@
 #define MARKER_CHAR 0xff
 #define STUFF_MARKER 0x00
 
-void init_bit_source(struct bit_source* source, input_func_t* input_func, void* userdata)
+void init_bit_source(struct bit_source* source, input_func_t* input_func,
+                     unsigned padding_bytes, void* userdata)
 {
 	memset(source, 0, sizeof(*source));
 	source->bytes = (uint8_t*)malloc(BYTERESERVOIR_SIZE);
 	source->byte_read_ptr = source->bytes;
 	source->input_func = input_func;
+	source->padding_bytes_available = padding_bytes;
 	source->userdata = userdata;
 }
 
@@ -44,11 +46,22 @@ void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits)
 			exit(1);
 		}
 		if (bytes_read == 0) {
-			fprintf(stderr, "Premature EOF\n");
-			exit(1);
+			source->source_eof = true;
+			if (source->padding_bytes_available > 0) {
+				unsigned padding_to_add = source->padding_bytes_available;
+				if (padding_to_add > bytes_to_read) {
+					padding_to_add = bytes_to_read;
+				}
+				memset(source->bytes + source->bytes_available, 0, padding_to_add);
+				source->padding_bytes_available -= padding_to_add;
+				source->bytes_available += padding_to_add;
+			} else {
+				fprintf(stderr, "Premature EOF\n");
+				exit(1);
+			}
+		} else {
+			source->bytes_available += bytes_read;
 		}
-		
-		source->bytes_available += bytes_read;
 	}
 
 	// Fill the bit reservoir one by one byte until we have enough.