From c0e4ca7a178330aad020899608c50fcf502052be Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Sun, 2 Apr 2017 14:15:33 +0200
Subject: [PATCH] Fix truncated gzip input

Since io.ErrUnexpectedEOF is returned on truncated results, we need to distinguish between truncated and end-of-stream reads.

Fixes https://github.com/klauspost/pgzip/issues/13
---
 gunzip.go      | 11 +++++++-
 gunzip_test.go | 71 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/gunzip.go b/gunzip.go
index 45b846e..f0e8fcb 100644
--- a/gunzip.go
+++ b/gunzip.go
@@ -388,7 +388,16 @@ func (z *Reader) doReadAhead() {
 			// Try to fill the buffer
 			n, err := io.ReadFull(decomp, buf)
 			if err == io.ErrUnexpectedEOF {
-				err = nil
+				if n > 0 {
+					err = nil
+				} else {
+					// If we got zero bytes, we need to establish if
+					// we reached end of stream or truncated stream.
+					_, err = decomp.Read([]byte{})
+					if err == io.EOF {
+						err = nil
+					}
+				}
 			}
 			if n < len(buf) {
 				buf = buf[0:n]
diff --git a/gunzip_test.go b/gunzip_test.go
index 1ebf56a..8e7d12e 100644
--- a/gunzip_test.go
+++ b/gunzip_test.go
@@ -596,3 +596,74 @@ func BenchmarkGunzipFlate(b *testing.B) {
 		}
 	}
 }
+
+func TestTruncatedGunzip(t *testing.T) {
+	in := []byte(strings.Repeat("ASDFASDFASDFASDFASDF", 1000))
+	var buf bytes.Buffer
+	enc := kpgzip.NewWriter(&buf)
+	_, err := enc.Write(in)
+	if err != nil {
+		t.Fatal(err)
+	}
+	enc.Close()
+	testdata := buf.Bytes()
+	for i := 5; i < len(testdata); i += 10 {
+		timer := time.NewTimer(time.Second)
+		done := make(chan struct{})
+		fail := make(chan struct{})
+		go func() {
+			r, err := NewReader(bytes.NewBuffer(testdata[:i]))
+			if err == nil {
+				b, err := ioutil.ReadAll(r)
+				if err == nil && !bytes.Equal(testdata[:i], b) {
+					close(fail)
+				}
+			}
+			close(done)
+		}()
+		select {
+		case <-timer.C:
+			t.Fatal("Timeout decoding")
+		case <-fail:
+			t.Fatal("No error, but mismatch")
+		case <-done:
+			timer.Stop()
+		}
+	}
+}
+
+func TestTruncatedGunzipBlocks(t *testing.T) {
+	var in = make([]byte, 512*10)
+	rand.Read(in)
+	var buf bytes.Buffer
+	for i := 0; i < len(in); i += 512 {
+		enc,_ := kpgzip.NewWriterLevel(&buf, 0)
+		_, err := enc.Write(in[:i])
+		if err != nil {
+			t.Fatal(err)
+		}
+		enc.Close()
+
+		timer := time.NewTimer(time.Second)
+		done := make(chan struct{})
+		fail := make(chan struct{})
+		go func() {
+			r, err := NewReaderN(&buf, 512, 10)
+			if err == nil {
+				b, err := ioutil.ReadAll(r)
+				if err == nil && !bytes.Equal(b, in[:i]) {
+					close(fail)
+				}
+			}
+			close(done)
+		}()
+		select {
+		case <-timer.C:
+			t.Fatal("Timeout decoding")
+		case <-fail:
+			t.Fatal("No error, but mismatch")
+		case <-done:
+			timer.Stop()
+		}
+	}
+}