From 4469f5446ad754d87afb340f59c213aec2dc4fb8 Mon Sep 17 00:00:00 2001
From: Russ Cox <rsc@golang.org>
Date: Thu, 16 Jul 2020 16:56:53 -0400
Subject: [PATCH] compress/flate: fix another deflate Reset inconsistency

While investigating #34121, fixed by CL 193605,
I discovered another case where Reset was not quite
resetting enough.

This specific case is not a problem in Reset itself but
rather that the Huffman bit writer in one code path
is using uninitialized memory left over from a previous
block, making the compression not choose the optimal
compression method.

Fixes #34121.

Change-Id: I29245b28214d924e382f91e2c56b4b8a9b7da13d
Reviewed-on: https://go-review.googlesource.com/c/go/+/243140
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
---
 src/compress/flate/deflate_test.go       | 56 +++++++++++++++++-------
 src/compress/flate/huffman_bit_writer.go |  1 +
 2 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/src/compress/flate/deflate_test.go b/src/compress/flate/deflate_test.go
index 4b1ed549a46b32..49a0345fd1d4bc 100644
--- a/src/compress/flate/deflate_test.go
+++ b/src/compress/flate/deflate_test.go
@@ -512,33 +512,57 @@ func TestWriterReset(t *testing.T) {
 			t.Errorf("level %d Writer not reset after Reset", level)
 		}
 	}
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriter(w, NoCompression) })
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriter(w, DefaultCompression) })
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriter(w, BestCompression) })
-	dict := []byte("we are the world")
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriterDict(w, NoCompression, dict) })
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriterDict(w, DefaultCompression, dict) })
-	testResetOutput(t, func(w io.Writer) (*Writer, error) { return NewWriterDict(w, BestCompression, dict) })
+
+	levels := []int{0, 1, 2, 5, 9}
+	for _, level := range levels {
+		t.Run(fmt.Sprint(level), func(t *testing.T) {
+			testResetOutput(t, level, nil)
+		})
+	}
+
+	t.Run("dict", func(t *testing.T) {
+		for _, level := range levels {
+			t.Run(fmt.Sprint(level), func(t *testing.T) {
+				testResetOutput(t, level, nil)
+			})
+		}
+	})
 }
 
-func testResetOutput(t *testing.T, newWriter func(w io.Writer) (*Writer, error)) {
+func testResetOutput(t *testing.T, level int, dict []byte) {
+	writeData := func(w *Writer) {
+		msg := []byte("now is the time for all good gophers")
+		w.Write(msg)
+		w.Flush()
+
+		hello := []byte("hello world")
+		for i := 0; i < 1024; i++ {
+			w.Write(hello)
+		}
+
+		fill := bytes.Repeat([]byte("x"), 65000)
+		w.Write(fill)
+	}
+
 	buf := new(bytes.Buffer)
-	w, err := newWriter(buf)
+	var w *Writer
+	var err error
+	if dict == nil {
+		w, err = NewWriter(buf, level)
+	} else {
+		w, err = NewWriterDict(buf, level, dict)
+	}
 	if err != nil {
 		t.Fatalf("NewWriter: %v", err)
 	}
-	b := []byte("hello world")
-	for i := 0; i < 1024; i++ {
-		w.Write(b)
-	}
+
+	writeData(w)
 	w.Close()
 	out1 := buf.Bytes()
 
 	buf2 := new(bytes.Buffer)
 	w.Reset(buf2)
-	for i := 0; i < 1024; i++ {
-		w.Write(b)
-	}
+	writeData(w)
 	w.Close()
 	out2 := buf2.Bytes()
 
diff --git a/src/compress/flate/huffman_bit_writer.go b/src/compress/flate/huffman_bit_writer.go
index 3e19061f8b57e5..f111f9f592c596 100644
--- a/src/compress/flate/huffman_bit_writer.go
+++ b/src/compress/flate/huffman_bit_writer.go
@@ -634,6 +634,7 @@ func (w *huffmanBitWriter) writeBlockHuff(eof bool, input []byte) {
 	w.literalFreq[endBlockMarker] = 1
 
 	const numLiterals = endBlockMarker + 1
+	w.offsetFreq[0] = 1
 	const numOffsets = 1
 
 	w.literalEncoding.generate(w.literalFreq, 15)