Fix gzip support in socket_listener with tcp sockets (#7446)

internal/content_coding.go

@@ -1,18 +1,78 @@
 package internal
 
 import (
+	"bufio"
 	"bytes"
 	"compress/gzip"
 	"errors"
 	"io"
 )
 
+// NewStreamContentDecoder returns a reader that will decode the stream
+// according to the encoding type.
+func NewStreamContentDecoder(encoding string, r io.Reader) (io.Reader, error) {
+	switch encoding {
+	case "gzip":
+		return NewGzipReader(r)
+	case "identity", "":
+		return r, nil
+	default:
+		return nil, errors.New("invalid value for content_encoding")
+	}
+}
+
+// GzipReader is similar to gzip.Reader but reads only a single gzip stream per read.
+type GzipReader struct {
+	r           io.Reader
+	z           *gzip.Reader
+	endOfStream bool
+}
+
+func NewGzipReader(r io.Reader) (io.Reader, error) {
+	// We need a read that implements ByteReader in order to line up the next
+	// stream.
+	br := bufio.NewReader(r)
+
+	// Reads the first gzip stream header.
+	z, err := gzip.NewReader(br)
+	if err != nil {
+		return nil, err
+	}
+
+	// Prevent future calls to Read from reading the following gzip header.
+	z.Multistream(false)
+
+	return &GzipReader{r: br, z: z}, nil
+}
+
+func (r *GzipReader) Read(b []byte) (int, error) {
+	if r.endOfStream {
+		// Reads the next gzip header and prepares for the next stream.
+		err := r.z.Reset(r.r)
+		if err != nil {
+			return 0, err
+		}
+		r.z.Multistream(false)
+		r.endOfStream = false
+	}
+
+	n, err := r.z.Read(b)
+
+	// Since multistream is disabled, io.EOF indicates the end of the gzip
+	// sequence.  On the next read we must read the next gzip header.
+	if err == io.EOF {
+		r.endOfStream = true
+		return n, nil
+	}
+	return n, err
+
+}
+
 // NewContentEncoder returns a ContentEncoder for the encoding type.
 func NewContentEncoder(encoding string) (ContentEncoder, error) {
 	switch encoding {
 	case "gzip":
 		return NewGzipEncoder()
-
 	case "identity", "":
 		return NewIdentityEncoder(), nil
 	default:

internal/content_coding_test.go

@@ -1,6 +1,8 @@
 package internal
 
 import (
+	"bytes"
+	"io/ioutil"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -56,3 +58,37 @@ func TestIdentityEncodeDecode(t *testing.T) {
 
 	require.Equal(t, "howdy", string(actual))
 }
+
+func TestStreamIdentityDecode(t *testing.T) {
+	var r bytes.Buffer
+	n, err := r.Write([]byte("howdy"))
+	require.NoError(t, err)
+	require.Equal(t, 5, n)
+
+	dec, err := NewStreamContentDecoder("identity", &r)
+	require.NoError(t, err)
+
+	data, err := ioutil.ReadAll(dec)
+	require.NoError(t, err)
+
+	require.Equal(t, []byte("howdy"), data)
+}
+
+func TestStreamGzipDecode(t *testing.T) {
+	enc, err := NewGzipEncoder()
+	require.NoError(t, err)
+	written, err := enc.Encode([]byte("howdy"))
+	require.NoError(t, err)
+
+	w := bytes.NewBuffer(written)
+
+	dec, err := NewStreamContentDecoder("gzip", w)
+	require.NoError(t, err)
+
+	b := make([]byte, 10)
+	n, err := dec.Read(b)
+	require.NoError(t, err)
+	require.Equal(t, 5, n)
+
+	require.Equal(t, []byte("howdy"), b[:n])
+}