pansy/internal/imagenorm/imagenorm_test.go

package imagenorm

import (
	"bytes"
	"encoding/binary"
	"hash/crc32"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"os"
	"testing"
)

// pngBytes and jpegBytes generate in-memory fixtures for the two formats Go can
// encode; heic/webp come from testdata (Go has no encoder for them).
func pngBytes(t *testing.T, w, h int) []byte {
	t.Helper()
	m := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := range h {
		for x := range w {
			m.Pix[m.PixOffset(x, y)+0] = uint8(x)
			m.Pix[m.PixOffset(x, y)+3] = 255
		}
	}
	var b bytes.Buffer
	if err := png.Encode(&b, m); err != nil {
		t.Fatalf("encode png: %v", err)
	}
	return b.Bytes()
}

func jpegBytes(t *testing.T, w, h int) []byte {
	t.Helper()
	m := image.NewRGBA(image.Rect(0, 0, w, h))
	var b bytes.Buffer
	if err := jpeg.Encode(&b, m, nil); err != nil {
		t.Fatalf("encode jpeg: %v", err)
	}
	return b.Bytes()
}

func readTestdata(t *testing.T, name string) []byte {
	t.Helper()
	b, err := os.ReadFile("testdata/" + name)
	if err != nil {
		t.Fatalf("read %s: %v", name, err)
	}
	return b
}

// TestNormalizeAllFormats is the load-bearing test: every format pansy claims to
// accept must round-trip to a valid JPEG. It exists specifically to catch a
// dropped blank import — the failure mode where the common format (PNG) breaks
// while the exotic one (HEIC) works, because someone deleted `_ "image/png"`.
func TestNormalizeAllFormats(t *testing.T) {
	cases := []struct {
		name       string
		input      []byte
		wantFormat string
	}{
		{"png", pngBytes(t, 120, 90), "png"},
		{"jpeg", jpegBytes(t, 120, 90), "jpeg"},
		{"heic", readTestdata(t, "sample.heic"), "heic"},
		{"webp", readTestdata(t, "sample.webp"), "webp"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			out, format, err := Normalize(bytes.NewReader(tc.input), Options{})
			if err != nil {
				t.Fatalf("Normalize(%s): %v", tc.name, err)
			}
			if format != tc.wantFormat {
				t.Errorf("format = %q, want %q", format, tc.wantFormat)
			}
			// The output must itself be a decodable JPEG.
			_, outFormat, err := image.Decode(bytes.NewReader(out))
			if err != nil {
				t.Fatalf("output isn't a valid image: %v", err)
			}
			if outFormat != "jpeg" {
				t.Errorf("output format = %q, want jpeg", outFormat)
			}
		})
	}
}

// TestNormalizeDownscales checks a large image is shrunk to fit MaxDim on its
// longest edge with aspect ratio preserved, and a small one is left alone.
func TestNormalizeDownscales(t *testing.T) {
	// A 2:1 image twice as wide as DefaultMaxDim → clamped to DefaultMaxDim on the
	// long edge with aspect preserved. Derived from the constant, not hard-coded,
	// so the test tracks the default rather than silently asserting a magic number.
	longEdge := DefaultMaxDim * 2
	big := pngBytes(t, longEdge, longEdge/2)
	out, _, err := Normalize(bytes.NewReader(big), Options{})
	if err != nil {
		t.Fatalf("Normalize: %v", err)
	}
	cfg, _, err := image.DecodeConfig(bytes.NewReader(out))
	if err != nil {
		t.Fatalf("decode out: %v", err)
	}
	if cfg.Width != DefaultMaxDim {
		t.Errorf("width = %d, want %d (longest edge clamped)", cfg.Width, DefaultMaxDim)
	}
	if cfg.Height != DefaultMaxDim/2 {
		t.Errorf("height = %d, want %d (aspect preserved)", cfg.Height, DefaultMaxDim/2)
	}

	// A small image within bounds keeps its dimensions.
	small := pngBytes(t, 100, 80)
	out, _, err = Normalize(bytes.NewReader(small), Options{})
	if err != nil {
		t.Fatalf("Normalize small: %v", err)
	}
	cfg, _, err = image.DecodeConfig(bytes.NewReader(out))
	if err != nil {
		t.Fatalf("decode small out: %v", err)
	}
	if cfg.Width != 100 || cfg.Height != 80 {
		t.Errorf("small image resized to %dx%d, want 100x80", cfg.Width, cfg.Height)
	}
}

// TestNormalizeRejectsOversizeInput: an input past the byte cap is ErrTooLarge,
// refused without a full decode.
func TestNormalizeRejectsOversizeInput(t *testing.T) {
	big := pngBytes(t, 500, 500)
	_, _, err := Normalize(bytes.NewReader(big), Options{MaxBytes: 100})
	if err != ErrTooLarge {
		t.Errorf("over-cap input err = %v, want ErrTooLarge", err)
	}
}

// TestNormalizeRejectsGarbage: unreadable-as-image bytes and a truncated image
// both fail cleanly with ErrUnsupported, not a panic.
func TestNormalizeRejectsGarbage(t *testing.T) {
	_, _, err := Normalize(bytes.NewReader([]byte("not an image at all")), Options{})
	if err != ErrUnsupported {
		t.Errorf("garbage err = %v, want ErrUnsupported", err)
	}
	// A truncated image (valid header, cut body) also fails cleanly, not a panic.
	png := pngBytes(t, 100, 100)
	_, _, err = Normalize(bytes.NewReader(png[:len(png)/2]), Options{})
	if err != ErrUnsupported {
		t.Errorf("truncated image err = %v, want ErrUnsupported", err)
	}
}

// pngHeader builds a valid PNG signature + IHDR chunk (with a correct CRC, which
// DecodeConfig verifies) for the given dimensions, and nothing else. It's enough
// for image.DecodeConfig to report width/height without a real bitmap — exactly
// what's needed to exercise the pre-decode size guard with a tiny input.
func pngHeader(w, h uint32) []byte {
	var buf bytes.Buffer
	buf.Write([]byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a})
	ihdr := make([]byte, 13)
	binary.BigEndian.PutUint32(ihdr[0:], w)
	binary.BigEndian.PutUint32(ihdr[4:], h)
	ihdr[8] = 8 // bit depth
	ihdr[9] = 6 // colour type: RGBA
	// compression/filter/interlace already 0.
	binary.Write(&buf, binary.BigEndian, uint32(len(ihdr)))
	chunk := append([]byte("IHDR"), ihdr...)
	buf.Write(chunk)
	binary.Write(&buf, binary.BigEndian, crc32.ChecksumIEEE(chunk))
	return buf.Bytes()
}

// TestNormalizeRejectsPixelBomb is the guard the review found untested: a small
// input (a bare ~40-byte PNG header) claiming an enormous canvas is refused with
// ErrTooLarge from DecodeConfig alone, before image.Decode allocates anything.
// Covers both the per-side maxDimension trip and the pixel-count trip — and, via
// the near-2^16-per-side case, that the count math doesn't overflow.
func TestNormalizeRejectsPixelBomb(t *testing.T) {
	cases := []struct {
		name string
		w, h uint32
	}{
		{"huge single side", 60000, 10},             // > maxDimension on width
		{"huge area within side cap", 40000, 40000}, // sides < cap, area 1.6 GP > maxDecodePixels
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			hdr := pngHeader(tc.w, tc.h)
			if len(hdr) > 100 {
				t.Fatalf("header unexpectedly large (%d bytes) — not a bomb test", len(hdr))
			}
			// Sanity: the header really does decode to those dimensions.
			cfg, _, err := image.DecodeConfig(bytes.NewReader(hdr))
			if err != nil {
				t.Fatalf("crafted PNG header didn't parse: %v", err)
			}
			if uint32(cfg.Width) != tc.w || uint32(cfg.Height) != tc.h {
				t.Fatalf("header reports %dx%d, want %dx%d", cfg.Width, cfg.Height, tc.w, tc.h)
			}
			if _, _, err := Normalize(bytes.NewReader(hdr), Options{}); err != ErrTooLarge {
				t.Errorf("pixel bomb %dx%d err = %v, want ErrTooLarge", tc.w, tc.h, err)
			}
		})
	}
}

// orientedJPEG builds a JPEG whose top-left quadrant is white and the rest black
// — a marker to track through a rotation — tagged with the given EXIF orientation
// (1..8). The marker lets a test assert the pixels actually moved to where that
// orientation says they should.
func orientedJPEG(t *testing.T, orient int) []byte {
	t.Helper()
	const w, h = 40, 24
	m := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := range h {
		for x := range w {
			c := color.RGBA{0, 0, 0, 255}
			if x < w/2 && y < h/2 {
				c = color.RGBA{255, 255, 255, 255}
			}
			m.Set(x, y, c)
		}
	}
	var jb bytes.Buffer
	if err := jpeg.Encode(&jb, m, &jpeg.Options{Quality: 95}); err != nil {
		t.Fatalf("encode jpeg: %v", err)
	}
	if orient == 0 {
		return jb.Bytes() // caller wants a plain JPEG with no EXIF
	}
	return spliceExifOrientation(t, jb.Bytes(), orient)
}

// spliceExifOrientation inserts a minimal little-endian Exif APP1 segment
// carrying just the Orientation tag right after the JPEG SOI marker.
func spliceExifOrientation(t *testing.T, jpg []byte, orient int) []byte {
	t.Helper()
	var tiff bytes.Buffer
	tiff.WriteString("II") // little-endian
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(0x2A))
	_ = binary.Write(&tiff, binary.LittleEndian, uint32(8)) // IFD0 offset
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(1)) // one entry
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(0x0112))
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(3))      // SHORT
	_ = binary.Write(&tiff, binary.LittleEndian, uint32(1))      // count
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(orient)) // value
	_ = binary.Write(&tiff, binary.LittleEndian, uint16(0))      // value pad
	_ = binary.Write(&tiff, binary.LittleEndian, uint32(0))      // next IFD
	payload := append([]byte("Exif\x00\x00"), tiff.Bytes()...)

	segLen := len(payload) + 2
	seg := []byte{0xFF, 0xE1, byte(segLen >> 8), byte(segLen)}
	seg = append(seg, payload...)

	out := make([]byte, 0, len(jpg)+len(seg))
	out = append(out, jpg[:2]...) // SOI
	out = append(out, seg...)
	return append(out, jpg[2:]...)
}

func bright(c color.Color) bool {
	r, g, b, _ := c.RGBA() // 16-bit
	return (r+g+b)/3 > 0x8000
}

// TestNormalizeAppliesExifOrientation is the #103 regression: a phone photo tagged
// "rotate 90°" must come out of Normalize with the pixels upright, not sideways —
// the re-encode strips EXIF, so the rotation has to be baked into the bitmap.
func TestNormalizeAppliesExifOrientation(t *testing.T) {
	// The white marker starts centred at (10,6) in the 40x24 source. For each
	// orientation, wantW/H is the corrected canvas and (mx,my) is where that
	// marker must land — derived from the same transform Normalize applies.
	cases := []struct {
		name         string
		orient       int
		wantW, wantH int
		mx, my       int
	}{
		{"none", 0, 40, 24, 10, 6},          // no EXIF → unchanged, marker top-left
		{"normal", 1, 40, 24, 10, 6},        // normal → unchanged
		{"mirror-h", 2, 40, 24, 29, 6},      // flip horizontal → top-right
		{"rotate-180", 3, 40, 24, 29, 17},   // → bottom-right
		{"mirror-v", 4, 40, 24, 10, 17},     // flip vertical → bottom-left
		{"transpose", 5, 24, 40, 6, 10},     // main diagonal (dims swap)
		{"rotate-90-cw", 6, 24, 40, 17, 10}, // → top-right (dims swap)
		{"transverse", 7, 24, 40, 17, 29},   // anti-diagonal (dims swap)
		{"rotate-90-ccw", 8, 24, 40, 6, 29}, // → bottom-left (dims swap)
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			out, format, err := Normalize(bytes.NewReader(orientedJPEG(t, tc.orient)), Options{})
			if err != nil {
				t.Fatalf("Normalize err = %v", err)
			}
			if format != "jpeg" {
				t.Errorf("format = %q, want jpeg", format)
			}
			m, _, err := image.Decode(bytes.NewReader(out))
			if err != nil {
				t.Fatalf("decode output: %v", err)
			}
			if m.Bounds().Dx() != tc.wantW || m.Bounds().Dy() != tc.wantH {
				t.Errorf("output %dx%d, want %dx%d",
					m.Bounds().Dx(), m.Bounds().Dy(), tc.wantW, tc.wantH)
			}
			if !bright(m.At(m.Bounds().Min.X+tc.mx, m.Bounds().Min.Y+tc.my)) {
				t.Errorf("white marker not at (%d,%d) — orientation not applied", tc.mx, tc.my)
			}
		})
	}
}

// TestExifOrientationParsing pins the parser against non-JPEG and no-EXIF inputs,
// which must default to 1 (never guess a rotation onto a correct image).
func TestExifOrientationParsing(t *testing.T) {
	if o := exifOrientation(pngBytes(t, 8, 8)); o != 1 {
		t.Errorf("PNG orientation = %d, want 1 (no JPEG EXIF path)", o)
	}
	if o := exifOrientation(orientedJPEG(t, 0)); o != 1 {
		t.Errorf("JPEG without EXIF orientation = %d, want 1", o)
	}
	if o := exifOrientation(orientedJPEG(t, 6)); o != 6 {
		t.Errorf("JPEG tagged 6 → %d, want 6", o)
	}
	if o := exifOrientation([]byte("not an image")); o != 1 {
		t.Errorf("garbage bytes → %d, want 1", o)
	}
}