vp8: limit all other partitions to 1<<24 bytes in size, not just N-1 of
them.

Also check for int32 overflow in the webp package.

Fixes #10790.

Change-Id: Id1162fad8a467a72a5379c7f4432d64ef25bc37a
Reviewed-on: https://go-review.googlesource.com/10072
Reviewed-by: Rob Pike <r@golang.org>

 vp8/decode.go       | 54 ++++++++++++++++++++++++++++++++++++++++-------------
 webp/decode.go      |  2 +-
 webp/decode_test.go | 19 +++++++++++++++++++
 3 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/vp8/decode.go b/vp8/decode.go
index 0aafacd..1bb5028 100644
--- a/vp8/decode.go
+++ b/vp8/decode.go
@@ -270,25 +270,53 @@ func (d *Decoder) parseFilterHeader() {
 
 // parseOtherPartitions parses the other partitions, as specified in section 9.5.
 func (d *Decoder) parseOtherPartitions() error {
-	buf := make([]byte, d.r.n)
-	if err := d.r.ReadFull(buf); err != nil {
-		return err
-	}
+	const maxNOP = 1 << 3
+	var partLens [maxNOP]int
 	d.nOP = 1 << d.fp.readUint(uniformProb, 2)
+
+	// The final partition length is implied by the the remaining chunk data
+	// (d.r.n) and the other d.nOP-1 partition lengths. Those d.nOP-1 partition
+	// lengths are stored as 24-bit uints, i.e. up to 16 MiB per partition.
 	n := 3 * (d.nOP - 1)
-	if n > len(buf) {
+	partLens[d.nOP-1] = d.r.n - n
+	if partLens[d.nOP-1] < 0 {
 		return io.ErrUnexpectedEOF
 	}
-	partLen, buf := buf[:n], buf[n:]
-	for i := 0; i < d.nOP-1; i++ {
-		m := int(partLen[3*i+0]) | int(partLen[3*i+1])<<8 | int(partLen[3*i+2])<<16
-		if m > len(buf) {
-			return io.ErrUnexpectedEOF
+	if n > 0 {
+		buf := make([]byte, n)
+		if err := d.r.ReadFull(buf); err != nil {
+			return err
+		}
+		for i := 0; i < d.nOP-1; i++ {
+			pl := int(buf[3*i+0]) | int(buf[3*i+1])<<8 | int(buf[3*i+2])<<16
+			if pl > partLens[d.nOP-1] {
+				return io.ErrUnexpectedEOF
+			}
+			partLens[i] = pl
+			partLens[d.nOP-1] -= pl
+		}
+	}
+
+	// We check if the final partition length can also fit into a 24-bit uint.
+	// Strictly speaking, this isn't part of the spec, but it guards against a
+	// malicious WEBP image that is too large to ReadFull the encoded DCT
+	// coefficients into memory, whether that's because the actual WEBP file is
+	// too large, or whether its RIFF metadata lists too large a chunk.
+	if 1<<24 <= partLens[d.nOP-1] {
+		return errors.New("vp8: too much data to decode")
+	}
+
+	buf := make([]byte, d.r.n)
+	if err := d.r.ReadFull(buf); err != nil {
+		return err
+	}
+	for i, pl := range partLens {
+		if i == d.nOP {
+			break
 		}
-		d.op[i].init(buf[:m])
-		buf = buf[m:]
+		d.op[i].init(buf[:pl])
+		buf = buf[pl:]
 	}
-	d.op[d.nOP-1].init(buf)
 	return nil
 }
 
diff --git a/webp/decode.go b/webp/decode.go
index 259bac6..60fb556 100644
--- a/webp/decode.go
+++ b/webp/decode.go
@@ -77,7 +77,7 @@ func decode(r io.Reader, configOnly bool) (image.Image, image.Config, error) {
 			unfilterAlpha(alpha, alphaStride, (buf[0]>>2)&0x03)
 
 		case fccVP8:
-			if wantAlpha {
+			if wantAlpha || int32(chunkLen) < 0 {
 				return nil, image.Config{}, errInvalidFormat
 			}
 			d := vp8.NewDecoder()
diff --git a/webp/decode_test.go b/webp/decode_test.go
index 581d29e..4b69f90 100644
--- a/webp/decode_test.go
+++ b/webp/decode_test.go
@@ -254,6 +254,25 @@ loop:
 	}
 }
 
+// TestDecodePartitionTooLarge tests that decoding a malformed WEBP image
+// doesn't try to allocate an unreasonable amount of memory. This WEBP image
+// claims a RIFF chunk length of 0x12345678 bytes (291 MiB) compressed,
+// independent of the actual image size (0 pixels wide * 0 pixels high).
+//
+// This is based on golang.org/issue/10790.
+func TestDecodePartitionTooLarge(t *testing.T) {
+	data := "RIFF\xff\xff\xff\x7fWEBPVP8 " +
+		"\x78\x56\x34\x12" + // RIFF chunk length.
+		"\xbd\x01\x00\x14\x00\x00\xb2\x34\x0a\x9d\x01\x2a\x96\x00\x67\x00"
+	_, err := Decode(strings.NewReader(data))
+	if err == nil {
+		t.Fatal("got nil error, want non-nil")
+	}
+	if got, want := err.Error(), "too much data"; !strings.Contains(got, want) {
+		t.Fatalf("got error %q, want something containing %q", got, want)
+	}
+}
+
 func benchmarkDecode(b *testing.B, filename string) {
 	data, err := ioutil.ReadFile("../testdata/blue-purple-pink-large." + filename + ".webp")
 	if err != nil {