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>

diff --git a/src/net/http/internal/chunked.go b/src/net/http/internal/chunked.go
index 2e62c00d5db5013e9e73c26cc59179b385f3e261..63f321d03b957dfcbacafc39acfc972d33e1eb4f 100644 (file)
--- 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 9abe1ab6d9d674da8be9d6636f5e75a548491e74..d06716591ab21f54f2805759ec9390a807d3b820 100644 (file)
--- 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)
+       }
+
+}