From 1af2f966341dd2daafd6f22730c643aa08c3bc3d Mon Sep 17 00:00:00 2001
From: Jamie Liu <jamieliu@google.com>
Date: Thu, 29 Feb 2024 10:37:50 -0800
Subject: [PATCH] Fix counting in compressio.Writer.Write().

compressio segments uncompressed data into chunks of `pool.chunkSize` each. In
each call to `Writer.Write()`, any part of the input that spans a full chunk is
passed without copying as an "inline buffer" to the worker pool, while parts of
the input that do not fill up a chunk are copied into temporary buffer
`pool.buf`, hereafter a "pooled buffer" since these are obtained from a global
`bufPool`. Pooled buffers are returned to `bufPool` once they are no longer
needed (i.e. once they have been compressed) to reduce heap allocations; doing
the same for inline buffers would be incorrect since their slices are still
owned by `Writer.Write()`'s caller.

`Writer.Write()` waits for compression of inline buffers to be completed (since
`io.Writer.Write(p)` is forbidden from retaining `p`), but does not wait for
compression of pooled buffers. Thus, at the beginning of any call to
`Writer.Write()`, there may be any number of pooled buffers queued for
compression, but no inline buffers. `Writer.Write()` uses this property to
distinguish between pooled and inline buffers without explicitly marking them
by counting existing pooled buffers as `pendingPre`, and queued inline buffers
as `pendingInline`.

The bug occurs when a call to `Writer.Write()` initially observes a non-empty
`pool.buf` (due to a previous call to `Write()`), and fills it completely with
at least `pool.chunkSize` more data to go (so `pool.buf`, followed by at least
one inline buffer, need to be enqueued). In this case, the enqueued pooled
buffer is not counted in `pendingPre` or `pendingInline`, so `Writer.Write()`
will incorrectly treat the pooled buffer as inline and forget about the last
inline buffer; thus the *following* call to `Writer.Write()` or
`Writer.Close()` will incorrectly treat the last inline buffer as pooled and
"return" it to `bufPool`, allowing the input slice's data to be corrupted by
reuse after the call to `Writer.Write()` has returned.

Fix this by counting such initial pooled buffers in `pendingPre`. Also
restructure the `callback` to hopefully make what is going on a little clearer.
(Unfortunately the panic does not catch this bug, since `Writer.Write()` stops
calling `pool.schedule()` once `pendingInline` reaches 0.)

PiperOrigin-RevId: 611521099
---
 pkg/compressio/compressio.go | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/pkg/compressio/compressio.go b/pkg/compressio/compressio.go
index f8e92fa96..dfd86869d 100644
--- a/pkg/compressio/compressio.go
+++ b/pkg/compressio/compressio.go
@@ -690,27 +690,30 @@ func (w *Writer) Write(p []byte) (int, error) {
 		pendingInline = 0
 	)
 	callback := func(c *chunk) error {
-		if pendingPre == 0 && pendingInline > 0 {
+		if pendingPre > 0 {
+			pendingPre--
+			err := w.flush(c)
+			c.uncompressed.Reset()
+			bufPool.Put(c.uncompressed)
+			return err
+		}
+		if pendingInline > 0 {
 			pendingInline--
 			return w.flush(c)
 		}
-		if pendingPre > 0 {
-			pendingPre--
-		}
-		err := w.flush(c)
-		c.uncompressed.Reset()
-		bufPool.Put(c.uncompressed)
-		return err
+		panic("both pendingPre and pendingInline exhausted")
 	}
 
 	for done := 0; done < len(p); {
 		// Construct an inline buffer if we're doing an inline
 		// encoding; see above regarding the bytes.MinRead constraint.
+		inline := false
 		if w.buf.Len() == 0 && len(p) >= done+int(w.chunkSize) && len(p) >= done+bytes.MinRead {
 			bufPool.Put(w.buf) // Return to the pool; never scheduled.
 			w.buf = bytes.NewBuffer(p[done : done+int(w.chunkSize)])
 			done += int(w.chunkSize)
 			pendingInline++
+			inline = true
 		}
 
 		// Do we need to flush w.buf? Note that this case should be hit
@@ -720,6 +723,9 @@ func (w *Writer) Write(p []byte) (int, error) {
 			if err := w.schedule(newChunk(nil, nil, nil, w.buf), callback); err != nil {
 				return done, err
 			}
+			if !inline {
+				pendingPre++
+			}
 			// Reset the buffer, since this has now been scheduled
 			// for compression. Note that this may be trampled
 			// immediately by the bufPool.Put(w.buf) above if the