dchain/media/scrub_test.go

package media

import (
	"bytes"
	"image"
	"image/color"
	"image/jpeg"
	"testing"
)

// TestScrubImageRemovesEXIF: our scrubber re-encodes via stdlib JPEG, which
// does not preserve EXIF by construction. We verify that a crafted input
// carrying an EXIF marker produces an output without one.
func TestScrubImageRemovesEXIF(t *testing.T) {
	// Build a JPEG that explicitly contains an APP1 EXIF segment.
	// Structure: JPEG SOI + APP1 with "Exif\x00\x00" header + real image data.
	var base bytes.Buffer
	img := image.NewRGBA(image.Rect(0, 0, 8, 8))
	for y := 0; y < 8; y++ {
		for x := 0; x < 8; x++ {
			img.Set(x, y, color.RGBA{uint8(x * 32), uint8(y * 32), 128, 255})
		}
	}
	if err := jpeg.Encode(&base, img, &jpeg.Options{Quality: 80}); err != nil {
		t.Fatalf("encode base: %v", err)
	}
	input := injectEXIF(t, base.Bytes())

	if !bytes.Contains(input, []byte("Exif\x00\x00")) {
		t.Fatalf("test setup broken: EXIF not injected")
	}
	// Also drop an identifiable string in the EXIF payload so we can prove
	// it's gone.
	if !bytes.Contains(input, []byte("SECRETGPS")) {
		t.Fatalf("test setup broken: EXIF marker not injected")
	}

	cleaned, mime, err := ScrubImage(input, "image/jpeg")
	if err != nil {
		t.Fatalf("ScrubImage: %v", err)
	}
	if mime != "image/jpeg" {
		t.Errorf("mime: got %q, want image/jpeg", mime)
	}
	// Verify the scrubbed output doesn't contain our canary string.
	if bytes.Contains(cleaned, []byte("SECRETGPS")) {
		t.Errorf("EXIF canary survived scrub — metadata not stripped")
	}
	// Verify the output doesn't contain the EXIF segment marker.
	if bytes.Contains(cleaned, []byte("Exif\x00\x00")) {
		t.Errorf("EXIF header string survived scrub")
	}
	// Output must still be a valid JPEG.
	if _, err := jpeg.Decode(bytes.NewReader(cleaned)); err != nil {
		t.Errorf("scrubbed output is not a valid JPEG: %v", err)
	}
}

// injectEXIF splices a synthetic APP1 EXIF segment after the JPEG SOI.
// Segment layout: FF E1 <len_hi> <len_lo> "Exif\0\0" + arbitrary payload.
// The payload is NOT valid TIFF — that's fine; stdlib JPEG decoder skips
// unknown APP1 segments rather than aborting.
func injectEXIF(t *testing.T, src []byte) []byte {
	t.Helper()
	if len(src) < 2 || src[0] != 0xFF || src[1] != 0xD8 {
		t.Fatalf("not a JPEG")
	}
	payload := []byte("Exif\x00\x00" + "SECRETGPS-51.5074N-0.1278W-Canon-EOS-R5")
	segmentLen := len(payload) + 2 // +2 = 2 bytes of len field itself
	var seg bytes.Buffer
	seg.Write([]byte{0xFF, 0xE1})
	seg.WriteByte(byte(segmentLen >> 8))
	seg.WriteByte(byte(segmentLen & 0xff))
	seg.Write(payload)
	out := make([]byte, 0, len(src)+seg.Len())
	out = append(out, src[:2]...) // SOI
	out = append(out, seg.Bytes()...)
	out = append(out, src[2:]...)
	return out
}

// TestScrubImageMIMEMismatch: rejects bytes that don't match claimed MIME.
func TestScrubImageMIMEMismatch(t *testing.T) {
	var buf bytes.Buffer
	img := image.NewRGBA(image.Rect(0, 0, 4, 4))
	jpeg.Encode(&buf, img, nil)
	// Claim it's a PNG.
	_, _, err := ScrubImage(buf.Bytes(), "image/png")
	if err == nil {
		t.Fatalf("expected ErrMIMEMismatch, got nil")
	}
}

// TestScrubImageDownscale: images over ImageMaxDim are shrunk.
func TestScrubImageDownscale(t *testing.T) {
	// Make a 2000×1000 image — larger dim 2000 > 1080.
	img := image.NewRGBA(image.Rect(0, 0, 2000, 1000))
	for y := 0; y < 1000; y++ {
		for x := 0; x < 2000; x++ {
			img.Set(x, y, color.RGBA{128, 64, 200, 255})
		}
	}
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 80}); err != nil {
		t.Fatalf("encode: %v", err)
	}
	cleaned, _, err := ScrubImage(buf.Bytes(), "image/jpeg")
	if err != nil {
		t.Fatalf("ScrubImage: %v", err)
	}
	decoded, err := jpeg.Decode(bytes.NewReader(cleaned))
	if err != nil {
		t.Fatalf("decode scrubbed: %v", err)
	}
	b := decoded.Bounds()
	if b.Dx() > ImageMaxDim || b.Dy() > ImageMaxDim {
		t.Errorf("not downscaled: got %dx%d, want max %d", b.Dx(), b.Dy(), ImageMaxDim)
	}
	// Aspect ratio roughly preserved (2:1 → 1080:540 with rounding slack).
	if b.Dx() != ImageMaxDim {
		t.Errorf("larger dim: got %d, want %d", b.Dx(), ImageMaxDim)
	}
}

// TestDetectMIME: a few magic-byte cases to ensure magic detection works.
func TestDetectMIME(t *testing.T) {
	cases := []struct {
		data []byte
		want string
	}{
		{[]byte("\xff\xd8\xff\xe0garbage"), "image/jpeg"},
		{[]byte("\x89PNG\r\n\x1a\n..."), "image/png"},
		{[]byte("GIF89a..."), "image/gif"},
		{[]byte{}, ""},
	}
	for _, tc := range cases {
		got := detectMIME(tc.data)
		if got != tc.want {
			t.Errorf("detectMIME(%q): got %q want %q", string(tc.data[:min(len(tc.data), 12)]), got, tc.want)
		}
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}