vp8l: avoid allocating many unused Huffman tree groups

A VP8L image contains a number of Huffman tree groups.
The number of groups is derived from the largest group index
observed in the file. The largest possible index is 2^16-1.

A VP8L file may included unused groups.

Decoding and storing all unused groups can cause excessive
memory allocation from a small input file.

Use the same approach to unused groups as libwebp:
When there are many groups (for some definition of "many"),
store the groups as a compacted array with unreferenced
groups removed, and rewrite all references to index into
the new array.

In addition, limit the total number of groups to 2600.
The libwebp encoder never produces more than this number of
groups, and there don't seem to be any other encoders in
common use.

Fixes golang/go#80069
Fixes CVE-2026-46603

Change-Id: I8738d53b67510a8c12b99fd4df55f4126a6a6964
Reviewed-on: https://go-review.googlesource.com/c/image/+/793460
Auto-Submit: Damien Neil <dneil@google.com>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
SLSA-Policy-Verified: SLSA Policy Verification Service <devtools-gerritcodereview-exitgate@google.com>
Reviewed-by: Nigel Tao <nigeltao@golang.org>
Reviewed-by: Neal Patel <nealpatel@google.com>
Reviewed-by: Nigel Tao <nigeltao@google.com>
diff --git a/testdata/gopher-doc.skip-hgroup.lossless.webp b/testdata/gopher-doc.skip-hgroup.lossless.webp
new file mode 100644
index 0000000..a28a5bb
--- /dev/null
+++ b/testdata/gopher-doc.skip-hgroup.lossless.webp
Binary files differ
diff --git a/testdata/large-huffman-index.lossless.webp.bz2 b/testdata/large-huffman-index.lossless.webp.bz2
new file mode 100644
index 0000000..17941de
--- /dev/null
+++ b/testdata/large-huffman-index.lossless.webp.bz2
Binary files differ
diff --git a/vp8l/decode.go b/vp8l/decode.go
index 4319487..0079485 100644
--- a/vp8l/decode.go
+++ b/vp8l/decode.go
@@ -225,6 +225,7 @@
 }
 
 // decodeHuffmanTree decodes a Huffman tree into h.
+// h may be nil, in which case we decode and discard the tree.
 func (d *decoder) decodeHuffmanTree(h *hTree, alphabetSize uint32) error {
 	useSimple, err := d.read(1)
 	if err != nil {
@@ -252,6 +253,9 @@
 				return err
 			}
 		}
+		if h == nil {
+			return nil
+		}
 		return h.buildSimple(nSymbols, symbols, alphabetSize)
 	}
 
@@ -274,6 +278,9 @@
 	if err = d.decodeCodeLengths(codeLengths, codeLengthCodeLengths[:]); err != nil {
 		return err
 	}
+	if h == nil {
+		return nil
+	}
 	return h.build(codeLengths)
 }
 
@@ -296,7 +303,23 @@
 func (d *decoder) decodeHuffmanGroups(w int32, h int32, topLevel bool, ccBits uint32) (
 	hGroups []hGroup, hPix []byte, hBits uint32, err error) {
 
+	// hPix maps tiles to hGroup indices.
+	// The number of hGroups is determined by the maximum index.
+	// VP8L permits unused hGroups in the input.
+	// While the number of used hGroups is bounded by the number of tiles,
+	// the number of decoded hGroups is only bounded by the maximum hGroup index (2^16-1).
+	//
+	// When the total number of hGroups is large, we store only the referenced hGroups
+	// and rewrite hGroup indices in hPix.
+	//
+	// When mapping is nil, we do no remapping.
+	// Otherwise, for hGroup i (where i is the real index of the group),
+	// mapping[i] is -1 for an unreferenced group or the remapped index otherwise.
+	var mapping []int16
+
 	maxHGroupIndex := 0
+	numHGroups := int16(1)
+
 	if topLevel {
 		useMeta, err := d.read(1)
 		if err != nil {
@@ -308,7 +331,10 @@
 				return nil, nil, 0, err
 			}
 			hBits += 2
-			hPix, err = d.decodePix(nTiles(w, hBits), nTiles(h, hBits), 0, false)
+			tileW := nTiles(w, hBits)
+			tileH := nTiles(h, hBits)
+			numTiles := tileH * tileW
+			hPix, err = d.decodePix(tileW, tileH, 0, false)
 			if err != nil {
 				return nil, nil, 0, err
 			}
@@ -318,15 +344,67 @@
 					maxHGroupIndex = i
 				}
 			}
+
+			// The standard allows for up to 2^16-1 hGroups.
+			//
+			// An input file containing a 1600x900 image can include
+			// and reference every possible group. With maximum-size trees
+			// (a ~12MiB input file), this is >3GiB of allocation.
+			//
+			// libwebp will decode a file containing 2^16-1 hGroups,
+			// but it won't produce one. It encodes at most 2600 hGroups
+			// (MAX_HUFF_IMAGE_SIZE, src/dec/common_dec.h)
+			// Reject any image using more.
+			const maxHuffImageSize = 2600
+			const _ int16 = maxHuffImageSize + 1 // must fit in int16
+			if maxHGroupIndex >= maxHuffImageSize {
+				return nil, nil, 0, errors.New("vp8l: too many Huffman trees")
+			}
+
+			// If maxHGroupIndex is too large, map indices to contiguous indices
+			// to avoid allocating large slices of unused Huffman trees.
+			//
+			// The definition of "too large" matches the one used in libwebp
+			// (see src/dec/vp8l_dec.c).
+			if maxHGroupIndex >= 1000 || int32(maxHGroupIndex) >= numTiles {
+				mapping = make([]int16, maxHGroupIndex+1)
+				for i := range mapping {
+					mapping[i] = -1
+				}
+				numHGroups = 0
+				for p := 0; p < len(hPix); p += 4 {
+					i := int(hPix[p])<<8 | int(hPix[p+1])
+					if mapping[i] == -1 {
+						mapping[i] = numHGroups
+						numHGroups++
+					}
+					hPix[p] = byte(mapping[i] >> 8)
+					hPix[p+1] = byte(mapping[i])
+				}
+			} else {
+				numHGroups = int16(maxHGroupIndex + 1)
+			}
 		}
 	}
-	hGroups = make([]hGroup, maxHGroupIndex+1)
-	for i := range hGroups {
+
+	hGroups = make([]hGroup, numHGroups)
+	for i := range maxHGroupIndex + 1 {
+		var hg *hGroup
+		if mapping == nil {
+			hg = &hGroups[i]
+		} else if mapping[i] != -1 {
+			hg = &hGroups[mapping[i]]
+		}
+
 		for j, alphabetSize := range alphabetSizes {
 			if j == 0 && ccBits > 0 {
 				alphabetSize += 1 << ccBits
 			}
-			if err := d.decodeHuffmanTree(&hGroups[i][j], alphabetSize); err != nil {
+			var hTreeDst *hTree
+			if hg != nil {
+				hTreeDst = &hg[j]
+			}
+			if err := d.decodeHuffmanTree(hTreeDst, alphabetSize); err != nil {
 				return nil, nil, 0, err
 			}
 		}
diff --git a/webp/decode_test.go b/webp/decode_test.go
index d6f4341..f0e6fce 100644
--- a/webp/decode_test.go
+++ b/webp/decode_test.go
@@ -6,9 +6,11 @@
 
 import (
 	"bytes"
+	"compress/bzip2"
 	"fmt"
 	"image"
 	"image/png"
+	"io"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -170,86 +172,123 @@
 }
 
 func TestDecodeVP8L(t *testing.T) {
-	testCases := []string{
-		"blue-purple-pink",
-		"blue-purple-pink-large",
-		"gopher-doc.1bpp",
-		"gopher-doc.2bpp",
-		"gopher-doc.4bpp",
-		"gopher-doc.8bpp",
-		"gopher-doc.with-alpha",
-		"tux",
-		"yellow_rose",
+	testCases := []struct {
+		name    string
+		f0      string
+		f1      string
+		wantErr string
+	}{
+		{name: "blue-purple-pink"},
+		{name: "blue-purple-pink-large"},
+		{name: "gopher-doc.1bpp"},
+		{name: "gopher-doc.2bpp"},
+		{name: "gopher-doc.4bpp"},
+		{name: "gopher-doc.8bpp"},
+		{name: "gopher-doc.with-alpha"},
+		{name: "tux"},
+		{name: "yellow_rose"},
+		{
+			// VP8L image with unreferenced Huffman tree groups.
+			name: "remapped hgroups",
+			f0:   "gopher-doc.skip-hgroup.lossless.webp",
+			f1:   "gopher-doc.8bpp.png",
+		},
+		{
+			// This file contains an image referencing Huffman tree group 65535,
+			// and trivial entries for all the preceding groups.
+			//
+			// When we allocated all groups (inculding unused ones), decoding this
+			// image allocated ~170MiB.
+			//
+			// We now reject this image for using more than 2600 hGroups.
+			name:    "large VP8L huffman index",
+			f0:      "large-huffman-index.lossless.webp.bz2",
+			wantErr: "vp8l: too many Huffman trees",
+		},
 	}
 
-loop:
+	openFile := func(t *testing.T, test, name, suffix string) io.Reader {
+		t.Helper()
+		if name == "" {
+			name = test + suffix
+		}
+		f, err := os.Open("../testdata/" + name)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() {
+			f.Close()
+		})
+		if strings.HasSuffix(name, ".bz2") {
+			return bzip2.NewReader(f)
+		}
+		return f
+	}
+
 	for _, tc := range testCases {
-		f0, err := os.Open("../testdata/" + tc + ".lossless.webp")
-		if err != nil {
-			t.Errorf("%s: Open WEBP: %v", tc, err)
-			continue
-		}
-		defer f0.Close()
-		img0, err := Decode(f0)
-		if err != nil {
-			t.Errorf("%s: Decode WEBP: %v", tc, err)
-			continue
-		}
-		m0, ok := img0.(*image.NRGBA)
-		if !ok {
-			t.Errorf("%s: WEBP image is %T, want *image.NRGBA", tc, img0)
-			continue
-		}
-
-		f1, err := os.Open("../testdata/" + tc + ".png")
-		if err != nil {
-			t.Errorf("%s: Open PNG: %v", tc, err)
-			continue
-		}
-		defer f1.Close()
-		img1, err := png.Decode(f1)
-		if err != nil {
-			t.Errorf("%s: Decode PNG: %v", tc, err)
-			continue
-		}
-		m1, ok := img1.(*image.NRGBA)
-		if !ok {
-			rgba1, ok := img1.(*image.RGBA)
+		t.Run(tc.name, func(t *testing.T) {
+			f0 := openFile(t, tc.name, tc.f0, ".lossless.webp")
+			img0, err := Decode(f0)
+			if tc.wantErr != "" {
+				if err == nil {
+					t.Fatalf("Decode WEBP: succeded, want error %q", tc.wantErr)
+				}
+				if !strings.Contains(err.Error(), tc.wantErr) {
+					t.Fatalf("Decode WEBP: error %q, want %q", err, tc.wantErr)
+				}
+				return
+			} else if err != nil {
+				t.Fatalf("Decode WEBP: error %q, want success", err)
+			}
+			m0, ok := img0.(*image.NRGBA)
 			if !ok {
-				t.Fatalf("%s: PNG image is %T, want *image.NRGBA", tc, img1)
-				continue
+				t.Fatalf("WEBP image is %T, want *image.NRGBA", img0)
 			}
-			if !rgba1.Opaque() {
-				t.Fatalf("%s: PNG image is non-opaque *image.RGBA, want *image.NRGBA", tc)
-				continue
-			}
-			// The image is fully opaque, so we can re-interpret the RGBA pixels
-			// as NRGBA pixels.
-			m1 = &image.NRGBA{
-				Pix:    rgba1.Pix,
-				Stride: rgba1.Stride,
-				Rect:   rgba1.Rect,
-			}
-		}
 
-		b0, b1 := m0.Bounds(), m1.Bounds()
-		if b0 != b1 {
-			t.Errorf("%s: bounds: got %v, want %v", tc, b0, b1)
-			continue
-		}
-		for i := range m0.Pix {
-			if m0.Pix[i] != m1.Pix[i] {
-				y := i / m0.Stride
-				x := (i - y*m0.Stride) / 4
-				i = 4 * (y*m0.Stride + x)
-				t.Errorf("%s: at (%d, %d):\ngot  %02x %02x %02x %02x\nwant %02x %02x %02x %02x",
-					tc, x, y,
-					m0.Pix[i+0], m0.Pix[i+1], m0.Pix[i+2], m0.Pix[i+3],
-					m1.Pix[i+0], m1.Pix[i+1], m1.Pix[i+2], m1.Pix[i+3],
-				)
-				continue loop
+			name1 := tc.f1
+			if name1 == "" {
+				name1 = tc.name + ".png"
 			}
-		}
+			f1 := openFile(t, tc.name, tc.f1, ".png")
+			img1, err := png.Decode(f1)
+			if err != nil {
+				t.Fatalf("Decode PNG: %v", err)
+			}
+			m1, ok := img1.(*image.NRGBA)
+			if !ok {
+				rgba1, ok := img1.(*image.RGBA)
+				if !ok {
+					t.Fatalf("PNG image is %T, want *image.NRGBA", img1)
+				}
+				if !rgba1.Opaque() {
+					t.Fatalf("PNG image is non-opaque *image.RGBA, want *image.NRGBA")
+				}
+				// The image is fully opaque, so we can re-interpret the RGBA pixels
+				// as NRGBA pixels.
+				m1 = &image.NRGBA{
+					Pix:    rgba1.Pix,
+					Stride: rgba1.Stride,
+					Rect:   rgba1.Rect,
+				}
+			}
+
+			b0, b1 := m0.Bounds(), m1.Bounds()
+			if b0 != b1 {
+				t.Fatalf("bounds: got %v, want %v", b0, b1)
+			}
+			for i := range m0.Pix {
+				if m0.Pix[i] != m1.Pix[i] {
+					y := i / m0.Stride
+					x := (i - y*m0.Stride) / 4
+					i = 4 * (y*m0.Stride + x)
+					t.Fatalf("at (%d, %d):\ngot  %02x %02x %02x %02x\nwant %02x %02x %02x %02x",
+						x, y,
+						m0.Pix[i+0], m0.Pix[i+1], m0.Pix[i+2], m0.Pix[i+3],
+						m1.Pix[i+0], m1.Pix[i+1], m1.Pix[i+2], m1.Pix[i+3],
+					)
+				}
+			}
+		})
 	}
 }