diff --git a/internal/media/image.go b/internal/media/image.go index b798c3a..905eccc 100644 --- a/internal/media/image.go +++ b/internal/media/image.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "image" - "image/draw" "image/gif" "image/jpeg" "image/png" @@ -25,6 +24,19 @@ const DefaultMaxDim = 896 // original bytes stay in the blob store untouched. const JPEGQuality = 85 +// DefaultMaxPixels — the largest source image this build will decode, counted +// in pixels rather than in compressed bytes. A byte cap is not a memory bound +// for an image: a 20000x20000 PNG of flat colour compresses to a few hundred +// kilobytes and decodes to 400 million pixels, which is 1.6 GB of heap in the +// process that owns the database and the socket. 40 megapixels is well past any +// phone camera and two orders of magnitude short of an OOM. +const DefaultMaxPixels = 40 << 20 + +// ErrTooManyPixels — the image header declares more pixels than this build +// will decode. Separate from ErrUnsupportedImage because the format is fine and +// the size is not, and the log line should say which. +var ErrTooManyPixels = errors.New("media: image has too many pixels") + // ErrUnsupportedImage — the bytes are not an image format this build can // decode. Notably webp: the stdlib has no webp decoder and this repo takes no // new dependencies, so a webp arriving from Telegram is refused here with a @@ -91,6 +103,16 @@ func PrepareImage(data []byte, source string, maxDim int) (Image, error) { if err != nil { return Image{}, err } + // The header is read before the pixels. Deciding after the decode is not a + // decision: by then the whole bitmap is already in the heap. + cfg, err := decodeConfig(data, mime) + if err != nil { + return Image{}, fmt.Errorf("media: read %s header: %w", mime, err) + } + if px := int64(cfg.Width) * int64(cfg.Height); px > DefaultMaxPixels { + return Image{}, fmt.Errorf("%w: %dx%d is %d, cap is %d", + ErrTooManyPixels, cfg.Width, cfg.Height, px, int64(DefaultMaxPixels)) + } src, err := decode(data, mime) if err != nil { return Image{}, fmt.Errorf("media: decode %s: %w", mime, err) @@ -105,6 +127,19 @@ func PrepareImage(data []byte, source string, maxDim int) (Image, error) { return Image{JPEG: buf.Bytes(), Width: b.Dx(), Height: b.Dy(), Source: source}, nil } +func decodeConfig(data []byte, mime string) (image.Config, error) { + r := bytes.NewReader(data) + switch strings.ToLower(mime) { + case "image/jpeg": + return jpeg.DecodeConfig(r) + case "image/png": + return png.DecodeConfig(r) + case "image/gif": + return gif.DecodeConfig(r) + } + return image.Config{}, ErrUnsupportedImage +} + func decode(data []byte, mime string) (image.Image, error) { r := bytes.NewReader(data) switch strings.ToLower(mime) { @@ -123,18 +158,16 @@ func decode(data []byte, mime string) (image.Image, error) { // destination pixel — nearest-neighbour would alias small text into noise, // which defeats the point of reading a screenshot, and an area average is a // dozen lines against pulling in golang.org/x/image on an offline box. +// +// It reads the source through At and allocates only the destination. Flattening +// into a full-size RGBA first doubled the peak: a 40-megapixel photo already +// costs 160 MB decoded, and the intermediate made it 320 MB before MaxDim had +// any chance to help. func flattenAndScale(src image.Image, maxDim int) *image.RGBA { sb := src.Bounds() sw, sh := sb.Dx(), sb.Dy() dw, dh := fit(sw, sh, maxDim) - flat := image.NewRGBA(image.Rect(0, 0, sw, sh)) - draw.Draw(flat, flat.Bounds(), image.NewUniform(image.White), image.Point{}, draw.Src) - draw.Draw(flat, flat.Bounds(), src, sb.Min, draw.Over) - if dw == sw && dh == sh { - return flat - } - dst := image.NewRGBA(image.Rect(0, 0, dw, dh)) for y := 0; y < dh; y++ { y0, y1 := y*sh/dh, (y+1)*sh/dh @@ -146,20 +179,24 @@ func flattenAndScale(src image.Image, maxDim int) *image.RGBA { if x1 <= x0 { x1 = x0 + 1 } - var r, g, b, n uint32 + var r, g, b, n uint64 for sy := y0; sy < y1; sy++ { for sx := x0; sx < x1; sx++ { - i := flat.PixOffset(sx, sy) - r += uint32(flat.Pix[i]) - g += uint32(flat.Pix[i+1]) - b += uint32(flat.Pix[i+2]) + // At returns premultiplied 16-bit. Compositing over white + // is then c + (1-alpha), which is the same answer the + // draw.Over pass used to give, one pixel at a time. + cr, cg, cb, ca := src.At(sb.Min.X+sx, sb.Min.Y+sy).RGBA() + inv := uint64(0xFFFF - ca) + r += uint64(cr) + inv + g += uint64(cg) + inv + b += uint64(cb) + inv n++ } } o := dst.PixOffset(x, y) - dst.Pix[o] = uint8(r / n) - dst.Pix[o+1] = uint8(g / n) - dst.Pix[o+2] = uint8(b / n) + dst.Pix[o] = uint8(r / n >> 8) + dst.Pix[o+1] = uint8(g / n >> 8) + dst.Pix[o+2] = uint8(b / n >> 8) dst.Pix[o+3] = 0xFF } } diff --git a/internal/media/image_test.go b/internal/media/image_test.go index 4355d0c..fe2802b 100644 --- a/internal/media/image_test.go +++ b/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() +}