kami
61ba58388f
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.
261 lines
7.6 KiB
Go
261 lines
7.6 KiB
Go
package media
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/binary"
|
||
"errors"
|
||
"hash/crc32"
|
||
"image"
|
||
"image/color"
|
||
"image/gif"
|
||
"image/jpeg"
|
||
"image/png"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
// pngBytes builds a w×h test image: left half red, right half a light grey, so
|
||
// a downscale that averages produces a predictable mid value and a scaler that
|
||
// silently returns the wrong region is visible.
|
||
func pngBytes(t *testing.T, w, h int) []byte {
|
||
t.Helper()
|
||
img := image.NewRGBA(image.Rect(0, 0, w, h))
|
||
for y := 0; y < h; y++ {
|
||
for x := 0; x < w; x++ {
|
||
if x < w/2 {
|
||
img.Set(x, y, color.RGBA{255, 0, 0, 255})
|
||
} else {
|
||
img.Set(x, y, color.RGBA{200, 200, 200, 255})
|
||
}
|
||
}
|
||
}
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, img); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return buf.Bytes()
|
||
}
|
||
|
||
func TestSniffImage(t *testing.T) {
|
||
cases := []struct {
|
||
name string
|
||
data []byte
|
||
want string
|
||
}{
|
||
{"png", pngBytes(t, 4, 4), "image/png"},
|
||
{"jpeg", jpegBytes(t, 4, 4), "image/jpeg"},
|
||
{"gif", gifBytes(t, 4, 4), "image/gif"},
|
||
}
|
||
for _, c := range cases {
|
||
got, err := SniffImage(c.data)
|
||
if err != nil {
|
||
t.Errorf("%s: %v", c.name, err)
|
||
continue
|
||
}
|
||
if got != c.want {
|
||
t.Errorf("%s: got %q want %q", c.name, got, c.want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// webp is common from Telegram and there is no stdlib decoder, so it must be
|
||
// refused by name rather than mis-sniffed or fed to a model as noise.
|
||
func TestSniffRefusesWebpByName(t *testing.T) {
|
||
webp := append([]byte("RIFF\x00\x00\x00\x00WEBP"), make([]byte, 8)...)
|
||
_, err := SniffImage(webp)
|
||
if !errors.Is(err, ErrUnsupportedImage) {
|
||
t.Fatalf("got %v, want ErrUnsupportedImage", err)
|
||
}
|
||
if !strings.Contains(err.Error(), "webp") {
|
||
t.Errorf("error does not name the format: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestSniffRefusesGarbage(t *testing.T) {
|
||
for _, data := range [][]byte{nil, []byte("hello"), []byte("\x00\x01\x02\x03")} {
|
||
if _, err := SniffImage(data); !errors.Is(err, ErrUnsupportedImage) {
|
||
t.Errorf("SniffImage(%q) = %v", data, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestPrepareImageDownscalesLongestEdge(t *testing.T) {
|
||
im, err := PrepareImage(pngBytes(t, 2000, 1000), "web:upload", 500)
|
||
if err != nil {
|
||
t.Fatalf("prepare: %v", err)
|
||
}
|
||
if im.Width != 500 || im.Height != 250 {
|
||
t.Errorf("got %dx%d, want 500x250", im.Width, im.Height)
|
||
}
|
||
if _, err := jpeg.Decode(bytes.NewReader(im.JPEG)); err != nil {
|
||
t.Errorf("output is not decodable jpeg: %v", err)
|
||
}
|
||
if im.Source != "web:upload" {
|
||
t.Errorf("source lost: %q", im.Source)
|
||
}
|
||
}
|
||
|
||
// Tall images scale on the other axis; a scaler that only handles landscape is
|
||
// the classic version of this bug.
|
||
func TestPrepareImageHandlesPortrait(t *testing.T) {
|
||
im, err := PrepareImage(pngBytes(t, 400, 1600), "telegram", 800)
|
||
if err != nil {
|
||
t.Fatalf("prepare: %v", err)
|
||
}
|
||
if im.Height != 800 || im.Width != 200 {
|
||
t.Errorf("got %dx%d, want 200x800", im.Width, im.Height)
|
||
}
|
||
}
|
||
|
||
func TestPrepareImageNeverEnlarges(t *testing.T) {
|
||
im, err := PrepareImage(pngBytes(t, 64, 32), "telegram", 896)
|
||
if err != nil {
|
||
t.Fatalf("prepare: %v", err)
|
||
}
|
||
if im.Width != 64 || im.Height != 32 {
|
||
t.Errorf("got %dx%d, want the original 64x32", im.Width, im.Height)
|
||
}
|
||
}
|
||
|
||
// A degenerate strip must not scale to zero on the short axis — jpeg.Encode
|
||
// fails on a zero-height image, which would turn a weird screenshot into a
|
||
// hard error.
|
||
func TestPrepareImageClampsDegenerateAspect(t *testing.T) {
|
||
im, err := PrepareImage(pngBytes(t, 2000, 2), "web:upload", 100)
|
||
if err != nil {
|
||
t.Fatalf("prepare: %v", err)
|
||
}
|
||
if im.Height < 1 || im.Width != 100 {
|
||
t.Errorf("got %dx%d", im.Width, im.Height)
|
||
}
|
||
}
|
||
|
||
// Transparent pixels composite onto white, not black: the common case is a
|
||
// screenshot or a diagram, and dark-on-black is unreadable to the model.
|
||
func TestPrepareImageFlattensAlphaOntoWhite(t *testing.T) {
|
||
img := image.NewRGBA(image.Rect(0, 0, 8, 8)) // fully transparent
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, img); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
im, err := PrepareImage(buf.Bytes(), "web:upload", 8)
|
||
if err != nil {
|
||
t.Fatalf("prepare: %v", err)
|
||
}
|
||
decoded, err := jpeg.Decode(bytes.NewReader(im.JPEG))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
r, g, b, _ := decoded.At(4, 4).RGBA()
|
||
if r>>8 < 240 || g>>8 < 240 || b>>8 < 240 {
|
||
t.Errorf("transparent pixel became rgb(%d,%d,%d), want near-white", r>>8, g>>8, b>>8)
|
||
}
|
||
}
|
||
|
||
func TestPrepareImageRejectsEmpty(t *testing.T) {
|
||
if _, err := PrepareImage(nil, "x", 0); !errors.Is(err, ErrEmpty) {
|
||
t.Errorf("got %v, want ErrEmpty", err)
|
||
}
|
||
}
|
||
|
||
func TestDataURIIsAJPEGDataURI(t *testing.T) {
|
||
im, err := PrepareImage(pngBytes(t, 16, 16), "x", 0)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
uri := im.DataURI()
|
||
if !strings.HasPrefix(uri, "data:image/jpeg;base64,") {
|
||
t.Fatalf("bad prefix: %.40s", uri)
|
||
}
|
||
if len(uri) <= len("data:image/jpeg;base64,") {
|
||
t.Error("data uri carries no payload")
|
||
}
|
||
}
|
||
|
||
func jpegBytes(t *testing.T, w, h int) []byte {
|
||
t.Helper()
|
||
img := image.NewRGBA(image.Rect(0, 0, w, h))
|
||
var buf bytes.Buffer
|
||
if err := jpeg.Encode(&buf, img, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return buf.Bytes()
|
||
}
|
||
|
||
func gifBytes(t *testing.T, w, h int) []byte {
|
||
t.Helper()
|
||
img := image.NewPaletted(image.Rect(0, 0, w, h), []color.Color{color.Black, color.White})
|
||
var buf bytes.Buffer
|
||
if err := gif.Encode(&buf, img, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
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()
|
||
}
|