From 73ee5b32c904b61f87c7fa693b86168db7957bc2 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 23 Jul 2023 11:17:16 +0530
Subject: [PATCH] A better way to wrap go's zlib into a streaming decompressor

---
 kittens/transfer/receive.go        |  58 +++--------------
 tools/utils/stream_decompressor.go | 101 +++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+), 50 deletions(-)
 create mode 100644 tools/utils/stream_decompressor.go

diff --git a/kittens/transfer/receive.go b/kittens/transfer/receive.go
index 337fe37c7..758e285e4 100644
--- a/kittens/transfer/receive.go
+++ b/kittens/transfer/receive.go
@@ -3,7 +3,6 @@
 package transfer
 
 import (
-	"bytes"
 	"compress/zlib"
 	"fmt"
 	"io"
@@ -29,48 +28,6 @@ const (
 	state_canceled
 )
 
-type decompressor = func([]byte, bool) ([]byte, error)
-
-func identity_decompressor(data []byte, is_last bool) ([]byte, error) { return data, nil }
-
-type zlib_decompressor struct {
-	z   io.ReadCloser
-	b   bytes.Buffer
-	buf []byte
-}
-
-func (self *zlib_decompressor) add_bytes(b []byte, is_last bool) (ans []byte, err error) {
-	self.b.Write(b)
-	pos, n := 0, 0
-	for {
-		if cap(self.buf) < pos+1024 {
-			newcap := utils.Max(2*cap(self.buf), pos+8192)
-			self.buf = append(self.buf[:pos], make([]byte, newcap-pos)...)
-		}
-		n, err = self.z.Read(self.buf[pos:cap(self.buf)])
-		pos += n
-		switch err {
-		case io.EOF:
-			n = 0
-		case nil:
-		default:
-			return nil, err
-		}
-		if n == 0 {
-			break
-		}
-	}
-	if is_last {
-		if err = self.z.Close(); err != nil {
-			return nil, err
-		}
-	}
-	if self.b.Len() == 0 {
-		self.b.Reset()
-	}
-	return self.buf[:pos], nil
-}
-
 type output_file interface {
 	write([]byte) error
 	close() error
@@ -115,7 +72,7 @@ type remote_file struct {
 	parent                       string
 	expanded_local_path          string
 	file_id                      string
-	decompressor                 decompressor
+	decompressor                 utils.StreamDecompressor
 	remote_symlink_value         string
 	actual_file                  output_file
 }
@@ -137,15 +94,16 @@ func new_remote_file(opts *Options, ftc *FileTransmissionCommand) (*remote_file,
 	ans := &remote_file{
 		expected_size: ftc.Size, ftype: ftc.Ftype, mtime: ftc.Mtime, spec_id: spec_id,
 		permissions: ftc.Permissions, remote_path: ftc.Name, display_name: wcswidth.StripEscapeCodes(ftc.Name),
-		remote_id: ftc.Status, remote_target: string(ftc.Data), parent: ftc.Parent, decompressor: identity_decompressor,
+		remote_id: ftc.Status, remote_target: string(ftc.Data), parent: ftc.Parent,
 	}
 	compression_capable := ftc.Ftype == FileType_regular && ftc.Size > 4096 && should_be_compressed(ftc.Name, opts.Compress)
 	if compression_capable {
-		z := &zlib_decompressor{}
-		if z.z, err = zlib.NewReader(&z.b); err != nil {
-			return nil, err
-		}
-		ans.decompressor = z.add_bytes
+		ans.decompressor, err = utils.NewStreamDecompressor(zlib.NewReader)
+	} else {
+		ans.decompressor, err = utils.NewStreamDecompressor(nil)
+	}
+	if err != nil {
+		return nil, err
 	}
 	return ans, nil
 }
diff --git a/tools/utils/stream_decompressor.go b/tools/utils/stream_decompressor.go
new file mode 100644
index 000000000..eda7e5d9a
--- /dev/null
+++ b/tools/utils/stream_decompressor.go
@@ -0,0 +1,101 @@
+// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package utils
+
+import (
+	"fmt"
+	"io"
+)
+
+var _ = fmt.Print
+
+// Decompress the data in chunk. output_callback() will be called zero or more times with decompressed data. Note that
+// it may be called in a different goroutine. The data provided to output_callback() is only valid for the lifetime of
+// output_callback(), so it should copy out the data.
+type StreamDecompressor = func(chunk []byte, is_last bool, output_callback func([]byte) error) error
+
+type output struct {
+	chunk []byte
+	err   error
+}
+
+type stream_decompressor struct {
+	impl     io.ReadCloser
+	pipe_r   *io.PipeReader
+	pipe_w   *io.PipeWriter
+	obuf     [8192]byte
+	err      error
+	callback func([]byte) error
+}
+
+func (self *stream_decompressor) process() {
+	for {
+		n, err := self.impl.Read(self.obuf[:])
+		if n > 0 {
+			if ocerr := self.callback(self.obuf[:n]); ocerr != nil {
+				self.pipe_r.CloseWithError(ocerr)
+				break
+			}
+		}
+		if err != nil {
+			self.pipe_r.CloseWithError(err)
+			break
+		}
+	}
+}
+
+func (self *stream_decompressor) next(chunk []byte, is_last bool, output_callback func([]byte) error) (err error) {
+	if self.err != nil {
+		return self.err
+	}
+	self.callback = output_callback
+	if _, err = self.pipe_w.Write(chunk); err != nil {
+		self.err = err
+		return err
+	}
+	if is_last {
+		defer func() {
+			self.pipe_r.Close()
+			self.pipe_w.Close()
+			if self.err == nil {
+				self.err = io.EOF
+			}
+		}()
+		self.err = self.impl.Close()
+		return self.err
+	}
+	return nil
+}
+
+// Wrap Go's awful decompressor routines to allow feeding them
+// data in chunks. For example:
+// sd, err := NewStreamDecompressor(zlib.NewReader)
+// sd(chunk, false, output_callback)
+// ...
+// sd(last_chunk, true, output_callback)
+// after this call calling sd() further will just return io.EOF
+func NewStreamDecompressor(constructor func(io.Reader) (io.ReadCloser, error)) (StreamDecompressor, error) {
+	if constructor == nil { // identity decompressor
+		var err error
+		return func(chunk []byte, is_last bool, cb func([]byte) error) error {
+			if err != nil {
+				return err
+			}
+			err = cb(chunk)
+			retval := err
+			if is_last && err != nil {
+				err = io.EOF
+			}
+			return retval
+		}, nil
+	}
+	s := stream_decompressor{}
+	s.pipe_r, s.pipe_w = io.Pipe()
+	rc, err := constructor(s.pipe_r)
+	if err != nil {
+		return nil, err
+	}
+	s.impl = rc
+	go s.process()
+	return s.next, nil
+}