From ee9a42282271861d405dbbadf35a014c5c6bf34c Mon Sep 17 00:00:00 2001
From: Lassi Pulkkinen <lassi@pulk.fi>
Date: Wed, 15 Mar 2023 07:46:15 +0200
Subject: [PATCH] Abort on attempt to create chunk reader at invalid position

To improve debuggability of bugs similar to the previous commit. Such a
bug could also occur in user code.

Signed-off-by: Lassi Pulkkinen <lassi@pulk.fi>
---
 image/png/reader.ha | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/image/png/reader.ha b/image/png/reader.ha
index 411f395..c3e6af1 100644
--- a/image/png/reader.ha
+++ b/image/png/reader.ha
@@ -7,6 +7,7 @@ use io;
 
 export type reader = struct {
 	src: io::handle,
+	nextchunk: bool,
 };
 
 // Creates a new PNG decoder. Reads and verifies the PNG magic before returning
@@ -24,6 +25,7 @@ export fn newreader(src: io::handle) (reader | error) = {
 	};
 	return reader {
 		src = src,
+		nextchunk = true,
 	};
 };
 
@@ -37,7 +39,7 @@ export fn newreader(src: io::handle) (reader | error) = {
 
 export type chunk_reader = struct {
 	vtable: io::stream,
-	src: io::handle,
+	src: *reader,
 	length: size,
 	ctype: u32,
 	crc: crc32::state,
@@ -52,6 +54,8 @@ export type chunk_reader = struct {
 // chunk, but instead will pass it to [[new_chunk_decoder]] to get a
 // chunk-type-aware decoder and read from this instead.
 export fn nextchunk(src: *reader) (chunk_reader | io::EOF | error) = {
+	assert(src.nextchunk,
+		"Must finish previous chunk before calling nextchunk again");
 	let buf: [8]u8 = [0...];
 	match (io::readall(src.src, buf[..])?) {
 	case io::EOF =>
@@ -59,13 +63,14 @@ export fn nextchunk(src: *reader) (chunk_reader | io::EOF | error) = {
 	case size =>
 		yield;
 	};
+	src.nextchunk = false;
 	const length = endian::begetu32(buf[..4]);
 	const ctype = endian::begetu32(buf[4..]);
 	const crc = crc32::crc32(&crc32::ieee_table);
 	hash::write(&crc, buf[4..]);
 	return chunk_reader {
 		vtable = &chunk_reader_vtable,
-		src = src.src,
+		src = src,
 		length = length,
 		ctype = ctype,
 		crc = crc,
@@ -134,12 +139,13 @@ fn chunk_read(st: *io::stream, buf: []u8) (size | io::EOF | io::error) = {
 
 	if (st.length == 0) {
 		let ckbuf: [4]u8 = [0...];
-		match (io::readall(st.src, ckbuf[..])?) {
+		match (io::readall(st.src.src, ckbuf[..])?) {
 		case io::EOF =>
 			return wraperror(invalid);
 		case size =>
 			yield;
 		};
+		st.src.nextchunk = true;
 		const want = endian::begetu32(ckbuf);
 		const have = crc32::sum32(&st.crc);
 		if (want != have) {
@@ -154,7 +160,7 @@ fn chunk_read(st: *io::stream, buf: []u8) (size | io::EOF | io::error) = {
 		yield st.length;
 	};
 
-	const z = match (io::read(st.src, buf[..max])?) {
+	const z = match (io::read(st.src.src, buf[..max])?) {
 	case io::EOF =>
 		// Missing checksum
 		return wraperror(invalid);
-- 
2.45.2