net/http/internal: don't block unnecessarily in ChunkedReader

Fixes #17355

Change-Id: I5390979cd0081b61a639466377faa46b4221b74a
Reviewed-on: https://go-review.googlesource.com/31329
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

2 files changed, 48 insertions, 9 deletions

diff --git a/src/net/http/internal/chunked.go b/src/net/http/internal/chunked.go
index 2e62c00d5d..63f321d03b 100644
--- a/src/net/http/internal/chunked.go
+++ b/src/net/http/internal/chunked.go
@@ -35,10 +35,11 @@ func NewChunkedReader(r io.Reader) io.Reader {
 }
 
 type chunkedReader struct {
-	r   *bufio.Reader
-	n   uint64 // unread bytes in chunk
-	err error
-	buf [2]byte
+	r        *bufio.Reader
+	n        uint64 // unread bytes in chunk
+	err      error
+	buf      [2]byte
+	checkEnd bool // whether need to check for \r\n chunk footer
 }
 
 func (cr *chunkedReader) beginChunk() {
@@ -68,6 +69,21 @@ func (cr *chunkedReader) chunkHeaderAvailable() bool {
 
 func (cr *chunkedReader) Read(b []uint8) (n int, err error) {
 	for cr.err == nil {
+		if cr.checkEnd {
+			if n > 0 && cr.r.Buffered() < 2 {
+				// We have some data. Return early (per the io.Reader
+				// contract) instead of potentially blocking while
+				// reading more.
+				break
+			}
+			if _, cr.err = io.ReadFull(cr.r, cr.buf[:2]); cr.err == nil {
+				if string(cr.buf[:]) != "\r\n" {
+					cr.err = errors.New("malformed chunked encoding")
+					break
+				}
+			}
+			cr.checkEnd = false
+		}
 		if cr.n == 0 {
 			if n > 0 && !cr.chunkHeaderAvailable() {
 				// We've read enough. Don't potentially block
@@ -92,11 +108,7 @@ func (cr *chunkedReader) Read(b []uint8) (n int, err error) {
 		// If we're at the end of a chunk, read the next two
 		// bytes to verify they are "\r\n".
 		if cr.n == 0 && cr.err == nil {
-			if _, cr.err = io.ReadFull(cr.r, cr.buf[:2]); cr.err == nil {
-				if cr.buf[0] != '\r' || cr.buf[1] != '\n' {
-					cr.err = errors.New("malformed chunked encoding")
-				}
-			}
+			cr.checkEnd = true
 		}
 	}
 	return n, cr.err
diff --git a/src/net/http/internal/chunked_test.go b/src/net/http/internal/chunked_test.go
index 9abe1ab6d9..d06716591a 100644
--- a/src/net/http/internal/chunked_test.go
+++ b/src/net/http/internal/chunked_test.go
@@ -185,3 +185,30 @@ func TestChunkReadingIgnoresExtensions(t *testing.T) {
 		t.Errorf("read %q; want %q", g, e)
 	}
 }
+
+// Issue 17355: ChunkedReader shouldn't block waiting for more data
+// if it can return something.
+func TestChunkReadPartial(t *testing.T) {
+	pr, pw := io.Pipe()
+	go func() {
+		pw.Write([]byte("7\r\n1234567"))
+	}()
+	cr := NewChunkedReader(pr)
+	readBuf := make([]byte, 7)
+	n, err := cr.Read(readBuf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := "1234567"
+	if n != 7 || string(readBuf) != want {
+		t.Fatalf("Read: %v %q; want %d, %q", n, readBuf[:n], len(want), want)
+	}
+	go func() {
+		pw.Write([]byte("xx"))
+	}()
+	_, err = cr.Read(readBuf)
+	if got := fmt.Sprint(err); !strings.Contains(got, "malformed") {
+		t.Fatalf("second read = %v; want malformed error", err)
+	}
+
+}