compress/gzip: allow stopping at end of first stream

Allows parsing some file formats that assign special
meaning to which stream data is found in.

Will do the same for compress/bzip2 once this is
reviewed and submitted.

Fixes #6486.

LGTM=nigeltao
R=nigeltao, dan.kortschak
CC=adg, bradfitz, golang-codereviews, r
https://golang.org/cl/159120044

2 files changed, 68 insertions, 0 deletions

diff --git a/src/compress/gzip/gunzip.go b/src/compress/gzip/gunzip.go
index df1d5aa2be..72ee55c4fa 100644
--- a/src/compress/gzip/gunzip.go
+++ b/src/compress/gzip/gunzip.go
@@ -74,6 +74,7 @@ type Reader struct {
 	flg          byte
 	buf          [512]byte
 	err          error
+	multistream  bool
 }
 
 // NewReader creates a new Reader reading the given reader.
@@ -83,6 +84,7 @@ type Reader struct {
 func NewReader(r io.Reader) (*Reader, error) {
 	z := new(Reader)
 	z.r = makeReader(r)
+	z.multistream = true
 	z.digest = crc32.NewIEEE()
 	if err := z.readHeader(true); err != nil {
 		return nil, err
@@ -102,9 +104,30 @@ func (z *Reader) Reset(r io.Reader) error {
 	}
 	z.size = 0
 	z.err = nil
+	z.multistream = true
 	return z.readHeader(true)
 }
 
+// Multistream controls whether the reader supports multistream files.
+//
+// If enabled (the default), the Reader expects the input to be a sequence
+// of individually gzipped data streams, each with its own header and
+// trailer, ending at EOF. The effect is that the concatenation of a sequence
+// of gzipped files is treated as equivalent to the gzip of the concatenation
+// of the sequence. This is standard behavior for gzip readers.
+//
+// Calling Multistream(false) disables this behavior; disabling the behavior
+// can be useful when reading file formats that distinguish individual gzip
+// data streams or mix gzip data streams with other data streams.
+// In this mode, when the Reader reaches the end of the data stream,
+// Read returns io.EOF. If the underlying reader implements io.ByteReader,
+// it will be left positioned just after the gzip stream.
+// To start the next stream, call z.Reset(r) followed by z.Multistream(false).
+// If there is no next stream, z.Reset(r) will return io.EOF.
+func (z *Reader) Multistream(ok bool) {
+	z.multistream = ok
+}
+
 // GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950).
 func get4(p []byte) uint32 {
 	return uint32(p[0]) | uint32(p[1])<<8 | uint32(p[2])<<16 | uint32(p[3])<<24
@@ -245,6 +268,10 @@ func (z *Reader) Read(p []byte) (n int, err error) {
 	}
 
 	// File is ok; is there another?
+	if !z.multistream {
+		return 0, io.EOF
+	}
+
 	if err = z.readHeader(false); err != nil {
 		z.err = err
 		return
diff --git a/src/compress/gzip/gunzip_test.go b/src/compress/gzip/gunzip_test.go
index 2471038f53..0636dec9ab 100644
--- a/src/compress/gzip/gunzip_test.go
+++ b/src/compress/gzip/gunzip_test.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"strings"
 	"testing"
 	"time"
 )
@@ -367,3 +368,43 @@ func TestInitialReset(t *testing.T) {
 		t.Errorf("got %q want %q", s, gunzipTests[1].raw)
 	}
 }
+
+func TestMultistreamFalse(t *testing.T) {
+	// Find concatenation test.
+	var tt gunzipTest
+	for _, tt = range gunzipTests {
+		if strings.HasSuffix(tt.desc, " x2") {
+			goto Found
+		}
+	}
+	t.Fatal("cannot find hello.txt x2 in gunzip tests")
+
+Found:
+	br := bytes.NewReader(tt.gzip)
+	var r Reader
+	if err := r.Reset(br); err != nil {
+		t.Fatalf("first reset: %v", err)
+	}
+
+	// Expect two streams with "hello world\n", then real EOF.
+	const hello = "hello world\n"
+
+	r.Multistream(false)
+	data, err := ioutil.ReadAll(&r)
+	if string(data) != hello || err != nil {
+		t.Fatalf("first stream = %q, %v, want %q, %v", string(data), err, hello, nil)
+	}
+
+	if err := r.Reset(br); err != nil {
+		t.Fatalf("second reset: %v", err)
+	}
+	r.Multistream(false)
+	data, err = ioutil.ReadAll(&r)
+	if string(data) != hello || err != nil {
+		t.Fatalf("second stream = %q, %v, want %q, %v", string(data), err, hello, nil)
+	}
+
+	if err := r.Reset(br); err != io.EOF {
+		t.Fatalf("third reset: err=%v, want io.EOF", err)
+	}
+}