path: root/apply_test.go

package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/gen2brain/go-fitz"
	"github.com/pdfcpu/pdfcpu/pkg/api"
)

// noTempLeftovers fails if any atomic-write temp file remains in dir.
func noTempLeftovers(t *testing.T, dir string) {
	t.Helper()
	leftovers, _ := filepath.Glob(filepath.Join(dir, ".*.tmp-*"))
	if len(leftovers) != 0 {
		t.Errorf("temp files left behind: %v", leftovers)
	}
}

// TestApply verifies end to end that a plan is burned into the output at the
// correct coordinates: a text stamp darkens its rect, an image stamp lands
// exactly in its rect (and not outside), and unstamped areas are untouched.
func TestApply(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src.pdf")
	imgPath := filepath.Join(dir, "stamp.png")
	out := filepath.Join(dir, "out.pdf")

	const ph = 792.0
	writeFile(t, src, blankPDF(2, "0 0 612 792", ""))
	writeFile(t, imgPath, redPNG(16, 16))

	p := &plan{
		Version: planVersion, Input: src,
		Stamps: []stamp{
			{ID: "t1", Type: typeText, Content: "TEST", Pages: []int{1},
				Rect: rect{X: 100, Y: 700, W: 120, H: 40}, Opacity: 1},
			{ID: "t2", Type: typeText, Content: strings.Repeat("W", 20), Pages: []int{1},
				Rect: rect{X: 100, Y: 300, W: 100, H: 30}, Opacity: 1},
			{ID: "i1", Type: typeImage, Src: imgPath, Pages: []int{1},
				Rect: rect{X: 400, Y: 100, W: 64, H: 64}, Opacity: 1},
		},
	}

	if err := apply(p, out); err != nil {
		t.Fatalf("apply: %v", err)
	}

	// The output must not keep the temp file's restrictive 0600 mode.
	if fi, err := os.Stat(out); err != nil {
		t.Fatalf("stat output: %v", err)
	} else if fi.Mode().Perm() != 0o644 {
		t.Errorf("output mode = %o, want 644", fi.Mode().Perm())
	}

	// Page count preserved.
	n, err := api.PageCountFile(out)
	if err != nil {
		t.Fatalf("PageCount: %v", err)
	}
	if n != 2 {
		t.Fatalf("page count = %d, want 2", n)
	}

	// render page 0 of both source and output at 72 dpi (1 px = 1 pt).
	srcImg := renderPage(t, src, 0)
	outImg := renderPage(t, out, 0)

	// pt (x,y) bottom-left -> px (x, ph-y) top-left at 72 dpi.
	px := func(x, y float64) (int, int) { return int(x), int(ph - y) }

	// Text rect darkened somewhere.
	if darkCount(outImg, 100, 60, 200, 90) == 0 {
		t.Errorf("text stamp produced no dark pixels in its rect")
	}

	// Wide text must fit its rect (100..200 x 300..330 pt): dark inside,
	// nothing spilling left or right.
	if darkCount(outImg, 100, 462, 200, 492) == 0 {
		t.Errorf("wide text stamp produced no dark pixels in its rect")
	}
	if n := darkCount(outImg, 40, 462, 96, 492); n != 0 {
		t.Errorf("wide text spilled left of its rect (%d dark px)", n)
	}
	if n := darkCount(outImg, 204, 462, 260, 492); n != 0 {
		t.Errorf("wide text spilled right of its rect (%d dark px)", n)
	}

	// Image stamp: center is red, registration is exact.
	cx, cy := px(432, 132) // center of 400..464 x 100..164
	if !isRed(outImg.At(cx, cy)) {
		t.Errorf("image stamp center (%d,%d) not red: %v", cx, cy, outImg.At(cx, cy))
	}
	// Just outside the image rect must remain white (no overflow / misplacement).
	ox, oy := px(480, 132)
	if !isWhite(outImg.At(ox, oy)) {
		t.Errorf("pixel right of image rect (%d,%d) not white: %v", ox, oy, outImg.At(ox, oy))
	}

	// Control: an unstamped area is unchanged white in both.
	mx, my := px(300, 400)
	if !isWhite(srcImg.At(mx, my)) || !isWhite(outImg.At(mx, my)) {
		t.Errorf("control pixel changed: src=%v out=%v", srcImg.At(mx, my), outImg.At(mx, my))
	}

	// Page 2 had no stamps: it must remain blank white.
	out2 := renderPage(t, out, 1)
	if darkCount(out2, 0, 0, 612, 792) != 0 {
		t.Errorf("page 2 should be blank but has dark pixels")
	}

	// The atomic write renamed its temp into place, leaving none behind.
	noTempLeftovers(t, dir)
}

// TestApplyRefusesInPlace verifies apply will not write over its own input,
// including when the output is a differently-spelled alias of the input.
func TestApplyRefusesInPlace(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "doc.pdf")
	writeFile(t, src, blankPDF(1, "0 0 612 792", ""))
	orig, _ := os.ReadFile(src)

	p := &plan{
		Version: planVersion, Input: src,
		Stamps: []stamp{{ID: "t", Type: typeText, Content: "X", Pages: []int{1},
			Rect: rect{X: 10, Y: 10, W: 50, H: 20}, Opacity: 1}},
	}

	if err := apply(p, src); err == nil || !strings.Contains(err.Error(), "same file") {
		t.Errorf("apply over input path: expected refusal, got %v", err)
	}

	link := filepath.Join(dir, "alias.pdf")
	if err := os.Symlink(src, link); err != nil {
		t.Skipf("symlink unsupported: %v", err)
	}
	if err := apply(p, link); err == nil || !strings.Contains(err.Error(), "same file") {
		t.Errorf("apply over symlink to input: expected refusal, got %v", err)
	}

	if got, _ := os.ReadFile(src); !bytes.Equal(got, orig) {
		t.Errorf("input was modified by a refused in-place apply")
	}
}

// TestApplyPreservesOutputOnFailure verifies a failed apply leaves an existing
// good output intact and removes its temp file.
func TestApplyPreservesOutputOnFailure(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src.pdf")
	out := filepath.Join(dir, "out.pdf")
	writeFile(t, src, blankPDF(1, "0 0 612 792", ""))

	sentinel := []byte("existing good output\n")
	writeFile(t, out, sentinel)

	// References a page beyond the document, so apply must fail.
	p := &plan{
		Version: planVersion, Input: src,
		Stamps: []stamp{{ID: "t", Type: typeText, Content: "X", Pages: []int{99},
			Rect: rect{X: 10, Y: 10, W: 50, H: 20}, Opacity: 1}},
	}
	if err := apply(p, out); err == nil {
		t.Fatal("expected apply to fail on out-of-range page")
	}

	if got, _ := os.ReadFile(out); !bytes.Equal(got, sentinel) {
		t.Errorf("existing output was modified on failure")
	}
	noTempLeftovers(t, dir)
}

func TestRejectsWebPContent(t *testing.T) {
	b, err := base64.StdEncoding.DecodeString("UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEAAUAmJQBOgCHwAP7+AA==")
	if err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(t.TempDir(), "stamp.png")
	writeFile(t, path, b)
	f, err := os.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	if _, _, err := decodeImageConfig(f); err == nil || !strings.Contains(err.Error(), "not supported") {
		t.Fatalf("expected WebP rejection, got %v", err)
	}
}

// TestApplyRotatedPages verifies stamps land at their picked (viewed-space)
// coordinates on pages with /Rotate: pdfcpu counter-rotates the overlay, so
// the red image must appear exactly at its rect in the rendered (viewed)
// output for every rotation.
func TestApplyRotatedPages(t *testing.T) {
	dir := t.TempDir()
	imgPath := filepath.Join(dir, "red.png")
	writeFile(t, imgPath, redPNG(16, 16))

	for _, rot := range []int{90, 180, 270} {
		src := filepath.Join(dir, fmt.Sprintf("src%d.pdf", rot))
		writeFile(t, src, blankPDF(1, "0 0 612 792", fmt.Sprintf("/Rotate %d", rot)))

		// Viewed height: a 612x792 page displays as 792x612 at 90/270.
		vh := 792
		if rot == 90 || rot == 270 {
			vh = 612
		}

		p := &plan{
			Version: planVersion, Input: src,
			Stamps: []stamp{{ID: "i", Type: typeImage, Src: imgPath, Pages: []int{1},
				Rect: rect{X: 100, Y: 100, W: 64, H: 64}, Opacity: 1}},
		}
		out := filepath.Join(dir, fmt.Sprintf("out%d.pdf", rot))
		if err := apply(p, out); err != nil {
			t.Fatalf("rot %d: apply: %v", rot, err)
		}

		img := renderPage(t, out, 0) // MuPDF renders the viewed orientation
		cx, cy := 132, vh-132        // center of 100..164 x 100..164, viewed
		if !isRed(img.At(cx, cy)) {
			t.Errorf("rot %d: stamp center (%d,%d) not red: %v", rot, cx, cy, img.At(cx, cy))
		}
		if ox, oy := 180, vh-132; !isWhite(img.At(ox, oy)) {
			t.Errorf("rot %d: pixel right of rect (%d,%d) not white: %v", rot, ox, oy, img.At(ox, oy))
		}
	}
}

// TestApplyRejectsInvalidRotate verifies a /Rotate that is not a multiple of
// 90 fails loudly.
func TestApplyRejectsInvalidRotate(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "r.pdf")
	writeFile(t, src, blankPDF(1, "0 0 612 792", "/Rotate 45"))

	p := &plan{
		Version: planVersion, Input: src,
		Stamps: []stamp{{ID: "t", Type: typeText, Content: "X", Pages: []int{1},
			Rect: rect{X: 10, Y: 10, W: 50, H: 20}, Opacity: 1}},
	}
	err := apply(p, filepath.Join(dir, "o.pdf"))
	if err == nil || !strings.Contains(err.Error(), "Rotate") {
		t.Fatalf("expected /Rotate rejection, got: %v", err)
	}
}

// TestApplyRejectsNonZeroOrigin verifies the page guard fails loudly on a
// page whose MediaBox origin is not (0,0) instead of misplacing stamps.
func TestApplyRejectsNonZeroOrigin(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "o.pdf")
	writeFile(t, src, blankPDF(1, "30 40 642 832", ""))

	p := &plan{
		Version: planVersion, Input: src,
		Stamps: []stamp{{ID: "t", Type: typeText, Content: "X", Pages: []int{1},
			Rect: rect{X: 10, Y: 10, W: 50, H: 20}, Opacity: 1}},
	}
	err := apply(p, filepath.Join(dir, "out.pdf"))
	if err == nil || !strings.Contains(err.Error(), "origin") {
		t.Fatalf("expected non-zero origin rejection, got: %v", err)
	}
}

// --- fixtures ---

// blankPDF builds a minimal valid PDF with n blank pages of the given
// MediaBox and extra page dict entries (e.g. "/Rotate 90").
func blankPDF(n int, mediaBox, extra string) []byte {
	pw := &pdfWriter{}
	cat := pw.reserve()
	pages := pw.reserve()
	kids := ""
	for i := 0; i < n; i++ {
		c := pw.add(makeStream("", []byte{}))
		page := pw.add([]byte(fmt.Sprintf(
			"<< /Type /Page /Parent %d 0 R /MediaBox [%s] %s /Resources << >> /Contents %d 0 R >>",
			pages, mediaBox, extra, c)))
		kids += fmt.Sprintf("%d 0 R ", page)
	}
	pw.set(pages, []byte(fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", kids, n)))
	pw.set(cat, []byte(fmt.Sprintf("<< /Type /Catalog /Pages %d 0 R >>", pages)))
	return pw.bytes()
}

func redPNG(w, h int) []byte {
	img := image.NewNRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			img.Set(x, y, color.NRGBA{R: 230, G: 20, B: 20, A: 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		panic(err)
	}
	return buf.Bytes()
}

// --- helpers ---

func writeFile(t *testing.T, path string, b []byte) {
	t.Helper()
	if err := os.WriteFile(path, b, 0o644); err != nil {
		t.Fatal(err)
	}
}

func renderPage(t *testing.T, path string, pg int) *image.RGBA {
	t.Helper()
	d, err := fitz.New(path)
	if err != nil {
		t.Fatalf("fitz open %s: %v", path, err)
	}
	defer d.Close()
	img, err := d.ImageDPI(pg, 72)
	if err != nil {
		t.Fatalf("render %s p%d: %v", path, pg, err)
	}
	return img
}

func darkCount(img *image.RGBA, x0, y0, x1, y1 int) int {
	n := 0
	b := img.Bounds()
	for y := y0; y < y1; y++ {
		for x := x0; x < x1; x++ {
			if x < b.Min.X || y < b.Min.Y || x >= b.Max.X || y >= b.Max.Y {
				continue
			}
			c := img.RGBAAt(x, y)
			if c.R < 100 && c.G < 100 && c.B < 100 {
				n++
			}
		}
	}
	return n
}

func isRed(c color.Color) bool {
	r, g, b, _ := c.RGBA()
	return r>>8 > 150 && g>>8 < 100 && b>>8 < 100
}

func isWhite(c color.Color) bool {
	r, g, b, _ := c.RGBA()
	return r>>8 > 240 && g>>8 > 240 && b>>8 > 240
}