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.
This commit is contained in:
kami
+53
-16
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user