media: bound an image by pixels, not by compressed bytes

The only cap was 64 MiB of input, and a decode bomb is a small file. A
20000x20000 PNG of flat colour compresses to a few hundred kilobytes,
decodes to 400 million pixels, and flattenAndScale then allocated a
second buffer of the same dimensions before scaling anything. That is
3.2 GB of live heap from one request, on a laptop, in the process that
owns the database and the socket, and max_dim never got a chance to
help. The header is read first now and a source over forty megapixels is
refused. The scaler reads the source through At and allocates only the
destination, so flattening no longer doubles the peak.

Found in review of #72.

internal/media/image_test.go

@@ -2,7 +2,9 @@ package media
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
+	"hash/crc32"
 	"image"
 	"image/color"
 	"image/gif"
@@ -189,3 +191,70 @@ func gifBytes(t *testing.T, w, h int) []byte {
 	}
 	return buf.Bytes()
 }
+
+// A decode bomb is a small file. Nothing bounded pixels before decoding, so a
+// 20000x20000 PNG of flat colour — a few hundred kilobytes on the wire, well
+// under the byte cap — decoded to 1.6 GB and then allocated another 1.6 GB to
+// flatten, in the process that owns the database and the socket.
+func TestPrepareImageRefusesADecodeBomb(t *testing.T) {
+	// The header is what is checked, so the test writes a real header and
+	// truncated pixel data: reaching the decode at all is the failure.
+	var buf bytes.Buffer
+	if err := png.Encode(&buf, image.NewGray(image.Rect(0, 0, 1, 1))); err != nil {
+		t.Fatal(err)
+	}
+	bomb := forgePNGSize(t, buf.Bytes(), 20000, 20000)
+	_, err := PrepareImage(bomb, "telegram", 0)
+	if !errors.Is(err, ErrTooManyPixels) {
+		t.Fatalf("err = %v, want ErrTooManyPixels", err)
+	}
+	// A phone photo is not a bomb.
+	if _, err := PrepareImage(pngBytes(t, 64, 48), "telegram", 0); err != nil {
+		t.Fatalf("an ordinary image was refused: %v", err)
+	}
+}
+
+// forgePNGSize rewrites the IHDR width and height (and its CRC) of a valid PNG,
+// which is how a header claiming 400 megapixels is produced without writing
+// 400 megapixels.
+func forgePNGSize(t *testing.T, src []byte, w, h uint32) []byte {
+	t.Helper()
+	out := append([]byte(nil), src...)
+	// 8 byte signature, 4 byte length, 4 byte "IHDR", then width and height.
+	const ihdr = 8 + 4 + 4
+	binary.BigEndian.PutUint32(out[ihdr:], w)
+	binary.BigEndian.PutUint32(out[ihdr+4:], h)
+	crc := crc32.ChecksumIEEE(out[8+4 : ihdr+13])
+	binary.BigEndian.PutUint32(out[ihdr+13:], crc)
+	return out
+}
+
+// Transparency still composites onto white, which is what makes a screenshot
+// readable. The old code did that with a full-size intermediate; the scaler
+// walks the source instead and must give the same answer.
+func TestPrepareImageFlattensOntoWhite(t *testing.T) {
+	img := image.NewRGBA(image.Rect(0, 0, 8, 8))
+	// Fully transparent everywhere: over white, that is white.
+	data := encodePNG(t, img)
+	out, err := PrepareImage(data, "test", 4)
+	if err != nil {
+		t.Fatal(err)
+	}
+	dec, err := jpeg.Decode(bytes.NewReader(out.JPEG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	r, g, b, _ := dec.At(2, 2).RGBA()
+	if r>>8 < 240 || g>>8 < 240 || b>>8 < 240 {
+		t.Fatalf("transparent pixel came out %d,%d,%d, want white", r>>8, g>>8, b>>8)
+	}
+}
+
+func encodePNG(t *testing.T, img image.Image) []byte {
+	t.Helper()
+	var buf bytes.Buffer
+	if err := png.Encode(&buf, img); err != nil {
+		t.Fatal(err)
+	}
+	return buf.Bytes()
+}