From 5f292889ec22c019a8630079874705b0fdaca0de Mon Sep 17 00:00:00 2001
From: Erik Corry <erik@toit.io>
Date: Fri, 4 Aug 2023 14:52:27 +0200
Subject: [PATCH] Make RLE encoder more robust. (#1726)

Allow double close and move the read method to a new reader class.
---
 lib/zlib.toit             | 64 ++++++++++++++++++++++++++++++---------
 tests/zlib_full_test.toit |  2 +-
 tests/zlib_test.toit      |  5 ++-
 3 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/lib/zlib.toit b/lib/zlib.toit
index da494d2fb..433179665 100644
--- a/lib/zlib.toit
+++ b/lib/zlib.toit
@@ -8,23 +8,49 @@ import crypto
 import crypto.adler32 as crypto
 import crypto.crc as crc_algorithms
 
+class RleReader implements reader.Reader:
+  owner_ /ZlibEncoder_? := null
+
+  constructor.private_:
+
+  read:
+    return owner_.read_
+
+  close:
+    // Currently does nothing.
+
+// TODO: Should not be a Reader, but we have deprecated the read method,
+// and there's no way to deprecate an 'implements' statement. This class
+// is a superclass for some public classes, so we can't just remove it
+// in a backwards compatible way.
 class ZlibEncoder_ implements reader.Reader:
   channel_ := monitor.Channel 1
+  reader/RleReader ::= RleReader.private_
 
   static SMALL_BUFFER_DEFLATE_HEADER_ ::= [8, 0x1d]
   static MINIMAL_GZIP_HEADER_ ::= [0x1f, 0x8b, 8, 0, 0, 0, 0, 0, 0, 0xff]
 
   constructor --gzip_header/bool:
+    reader.owner_ = this
     header := gzip_header ? MINIMAL_GZIP_HEADER_ : SMALL_BUFFER_DEFLATE_HEADER_
     channel_.send (ByteArray header.size: header[it])
 
+  /**
+  Deprecated.
+  Use the $reader to get an object with a read method.
+  */
   read:
-    return channel_.receive
+    return read_
+
+  read_:
+    result := channel_.receive
+    if not result: channel_.send null  // Once it's closed we should be able to keep reading nulls.
+    return result
 
 class UncompressedDeflateEncoder_ extends ZlibEncoder_:
-  summer_/crypto.Checksum := ?
+  summer_/crypto.Checksum? := ?
 
-  constructor this.summer_ --gzip_header/bool:
+  constructor .summer_ --gzip_header/bool:
     super --gzip_header=gzip_header
     buffer_fullness_ = BLOCK_HEADER_SIZE_
 
@@ -35,6 +61,7 @@ class UncompressedDeflateEncoder_ extends ZlibEncoder_:
   static BLOCK_HEADER_SIZE_ ::= 5
 
   write collection from=0 to=collection.size -> int:
+    if not summer_: throw "ALREADY_CLOSED"
     return List.chunk_up from to (buffer_.size - buffer_fullness_) buffer_.size: | from to bytes |
       buffer_.replace buffer_fullness_ collection from to
       buffer_fullness_ += bytes
@@ -59,10 +86,12 @@ class UncompressedDeflateEncoder_ extends ZlibEncoder_:
     buffer_fullness_ = 0
 
   close:
-    send_ --last=true
-    checksum_bytes := summer_.get
-    channel_.send checksum_bytes
-    channel_.send null
+    if summer_:
+      send_ --last=true
+      checksum_bytes := summer_.get
+      summer_ = null
+      channel_.send checksum_bytes
+      channel_.send null
 
 /**
 Creates an uncompressed data stream that is compatible with zlib decoders
@@ -124,6 +153,7 @@ class RunLengthDeflateEncoder_ extends ZlibEncoder_:
   write collection from=0 to=collection.size:
     summer_.add collection from to
     while to != from:
+      if not rle_: throw "ALREADY_CLOSED"
       // The buffer is 256 large, and we don't let it get too full because then the compressor
       // may not be able to make progress, so we flush it when we hit three quarters full.
       assert: buffer_.size == 256
@@ -140,14 +170,18 @@ class RunLengthDeflateEncoder_ extends ZlibEncoder_:
   Closes the encoder for writing.
   */
   close:
-    channel_.send (buffer_.copy 0 buffer_fullness_)
-    try:
-      written := rle_finish_ rle_ buffer_ 0
-      channel_.send (buffer_.copy 0 written)
-      channel_.send summer_.get
-      channel_.send null
-    finally:
-      remove_finalizer this
+    if buffer_fullness_ != 0:
+      channel_.send (buffer_.copy 0 buffer_fullness_)
+      buffer_fullness_ = 0
+    if rle_:
+      try:
+        written := rle_finish_ rle_ buffer_ 0
+        rle_ = null
+        channel_.send (buffer_.copy 0 written)
+        channel_.send summer_.get
+        channel_.send null
+      finally:
+        remove_finalizer this
 
 /**
 Creates a run length encoded data stream that is compatible with zlib decoders
diff --git a/tests/zlib_full_test.toit b/tests/zlib_full_test.toit
index 1c16890c4..86fb9e522 100644
--- a/tests/zlib_full_test.toit
+++ b/tests/zlib_full_test.toit
@@ -115,7 +115,7 @@ rle_test -> none:
     encoder.close
 
   encoded := #[]
-  while data := encoder.read:
+  while data := encoder.reader.read:
     encoded += data
 
   decoder := zlib.Decoder
diff --git a/tests/zlib_test.toit b/tests/zlib_test.toit
index 7d634a85e..025bda48f 100644
--- a/tests/zlib_test.toit
+++ b/tests/zlib_test.toit
@@ -26,7 +26,7 @@ do_test chunk_size chunk_offset str zlib_expected uncompressed --gzip/bool:
   accumulator := bytes.Buffer
   done := Semaphore
   t := task::
-    while ba := compressor.read:
+    while ba := compressor.reader.read:
       accumulator.write ba
     done.up
   if chunk_offset != 0:
@@ -36,6 +36,9 @@ do_test chunk_size chunk_offset str zlib_expected uncompressed --gzip/bool:
     compressor.write
       str.copy from to
   compressor.close
+  compressor.close  // Test it is idempotent.
+  e := catch: compressor.write "After close"
+  expect e != null
   done.down
   if not gzip:
     // Test output against expected compressed data.