4 files changed, 19 insertions, 17 deletions
diff --git a/reftable/block.c b/reftable/block.c
index 471faa1642..920b3f4486 100644
--- a/reftable/block.c
+++ b/reftable/block.c
@@ -227,7 +227,8 @@ static int read_block(struct reftable_block_source *source,
 int reftable_block_init(struct reftable_block *block,
 			struct reftable_block_source *source,
 			uint32_t offset, uint32_t header_size,
-			uint32_t table_block_size, uint32_t hash_size)
+			uint32_t table_block_size, uint32_t hash_size,
+			uint8_t want_type)
 {
 	uint32_t guess_block_size = table_block_size ?
 		table_block_size : DEFAULT_BLOCK_SIZE;
@@ -247,6 +248,10 @@ int reftable_block_init(struct reftable_block *block,
 		err = REFTABLE_FORMAT_ERROR;
 		goto done;
 	}
+	if (want_type != REFTABLE_BLOCK_TYPE_ANY && block_type != want_type) {
+		err = 1;
+		goto done;
+	}
 
 	block_size = reftable_get_be24(block->block_data.data + header_size + 1);
 	if (block_size > guess_block_size) {
diff --git a/reftable/reftable-block.h b/reftable/reftable-block.h
index 04c3b518c8..0b05a8f7e3 100644
--- a/reftable/reftable-block.h
+++ b/reftable/reftable-block.h
@@ -56,7 +56,8 @@ struct reftable_block {
 int reftable_block_init(struct reftable_block *b,
 			struct reftable_block_source *source,
 			uint32_t offset, uint32_t header_size,
-			uint32_t table_block_size, uint32_t hash_size);
+			uint32_t table_block_size, uint32_t hash_size,
+			uint8_t want_type);
 
 /* Release resources allocated by the block. */
 void reftable_block_release(struct reftable_block *b);
diff --git a/reftable/table.c b/reftable/table.c
index ee83127615..56362df0ed 100644
--- a/reftable/table.c
+++ b/reftable/table.c
@@ -173,16 +173,7 @@ int table_init_block(struct reftable_table *t, struct reftable_block *block,
 		return 1;
 
 	err = reftable_block_init(block, &t->source, next_off, header_off,
-				  t->block_size, hash_size(t->hash_id));
-	if (err < 0)
-		goto done;
-
-	if (want_typ != REFTABLE_BLOCK_TYPE_ANY && block->block_type != want_typ) {
-		err = 1;
-		goto done;
-	}
-
-done:
+				  t->block_size, hash_size(t->hash_id), want_typ);
 	if (err)
 		reftable_block_release(block);
 	return err;
diff --git a/t/unit-tests/t-reftable-block.c b/t/unit-tests/t-reftable-block.c
index 7dbd93601c..52f1dae1c9 100644
--- a/t/unit-tests/t-reftable-block.c
+++ b/t/unit-tests/t-reftable-block.c
@@ -64,7 +64,8 @@ static void t_ref_block_read_write(void)
 	block_writer_release(&bw);
 
 	block_source_from_buf(&source ,&block_data);
-	reftable_block_init(&block, &source, 0, header_off, block_size, REFTABLE_HASH_SIZE_SHA1);
+	reftable_block_init(&block, &source, 0, header_off, block_size,
+			    REFTABLE_HASH_SIZE_SHA1, REFTABLE_BLOCK_TYPE_REF);
 
 	block_iter_init(&it, &block);
 
@@ -153,7 +154,8 @@ static void t_log_block_read_write(void)
 	block_writer_release(&bw);
 
 	block_source_from_buf(&source, &block_data);
-	reftable_block_init(&block, &source, 0, header_off, block_size, REFTABLE_HASH_SIZE_SHA1);
+	reftable_block_init(&block, &source, 0, header_off, block_size,
+			    REFTABLE_HASH_SIZE_SHA1, REFTABLE_BLOCK_TYPE_LOG);
 
 	block_iter_init(&it, &block);
 
@@ -245,7 +247,8 @@ static void t_obj_block_read_write(void)
 	block_writer_release(&bw);
 
 	block_source_from_buf(&source, &block_data);
-	reftable_block_init(&block, &source, 0, header_off, block_size, REFTABLE_HASH_SIZE_SHA1);
+	reftable_block_init(&block, &source, 0, header_off, block_size,
+			    REFTABLE_HASH_SIZE_SHA1, REFTABLE_BLOCK_TYPE_OBJ);
 
 	block_iter_init(&it, &block);
 
@@ -329,7 +332,8 @@ static void t_index_block_read_write(void)
 	block_writer_release(&bw);
 
 	block_source_from_buf(&source, &block_data);
-	reftable_block_init(&block, &source, 0, header_off, block_size, REFTABLE_HASH_SIZE_SHA1);
+	reftable_block_init(&block, &source, 0, header_off, block_size,
+			    REFTABLE_HASH_SIZE_SHA1, REFTABLE_BLOCK_TYPE_INDEX);
 
 	block_iter_init(&it, &block);
 
@@ -411,7 +415,8 @@ static void t_block_iterator(void)
 	check_int(err, >, 0);
 
 	block_source_from_buf(&source, &data);
-	reftable_block_init(&block, &source, 0, 0, data.len, REFTABLE_HASH_SIZE_SHA1);
+	reftable_block_init(&block, &source, 0, 0, data.len,
+			    REFTABLE_HASH_SIZE_SHA1, REFTABLE_BLOCK_TYPE_REF);
 
 	err = reftable_block_init_iterator(&block, &it);
 	check_int(err, ==, 0);