From c4b62c2feae238ac8cb3280be604754aa4e20832 Mon Sep 17 00:00:00 2001
From: JustAnotherArchivist <JustAnotherArchivist@users.noreply.github.com>
Date: Tue, 27 Jul 2021 05:16:00 +0000
Subject: [PATCH] Fix piping when reads return less data than expected

---
 zstdwarccat | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/zstdwarccat b/zstdwarccat
index b65eeaf..6f92414 100755
--- a/zstdwarccat
+++ b/zstdwarccat
@@ -15,7 +15,16 @@ def get_dict(fp):
 	dictSize = struct.unpack('<I', dictSize)[0]
 	assert dictSize >= 4, 'dict too small'
 	assert dictSize < 100 * 1024**2, 'dict too large'
-	d = fp.read(dictSize)
+	ds = []
+	dlen = 0
+	while dlen < dictSize:
+		c = fp.read(dictSize - dlen)
+		if c is None or c == b'': # EOF
+			break
+		ds.append(c)
+		dlen += len(c)
+	d = b''.join(ds)
+	assert len(d) == dictSize, f'could not read dict fully: expected {dictSize}, got {len(d)}'
 	assert d.startswith(b'\x28\xB5\x2F\xFD') or d.startswith(b'\x37\xA4\x30\xEC'), 'not a valid dict'
 	if d.startswith(b'\x28\xB5\x2F\xFD'): # Compressed dict
 		# Decompress with unzstd