pansy/internal/imagenorm/imagenorm.go

// Package imagenorm normalizes an uploaded image to a JPEG the rest of pansy
// (and majordomo's vision path) can rely on, decoding the formats a phone
// actually produces.
//
// # Why this exists
//
// majordomo's own media pipeline is stdlib-based, so it cannot decode HEIC or
// WebP — and HEIC is the iPhone camera default. The seed-packet feature's very
// first input is "a photo from my phone", so without this the feature fails on
// the exact device that motivates it. Normalizing at the upload boundary means
// everything downstream only ever sees JPEG.
//
// # The CGO constraint
//
// pansy is CGO_ENABLED=0 (a single static binary — the reason modernc/sqlite was
// chosen over the C one), so a libheif *binding* is out. github.com/gen2brain/heic
// runs libheif as WebAssembly via wazero: pure Go, no cgo, and it registers with
// image.Decode like any other format. golang.org/x/image/webp is pure Go too.
//
// # Import-driven registry
//
// Go's image decoders register via blank imports, and the failure mode is
// backwards from intuition: forget "image/png" and PNG uploads fail with
// "unknown format" while the exotic HEIC still works. So the blank imports below
// are load-bearing, and TestNormalizeAllFormats exercises all four formats to
// keep them so.
package imagenorm

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"io"

	"golang.org/x/image/draw"

	// Decoders, registered with image.Decode by side effect. All four matter:
	// jpeg/png are the common cases, heic is the iPhone default, webp is common
	// on the web. Dropping any one silently breaks that format's uploads.
	_ "image/jpeg"
	_ "image/png"

	_ "github.com/gen2brain/heic"
	_ "golang.org/x/image/webp"
)

// Defaults chosen for the seed-packet path against ollama-cloud's limits (8
// images, 20 MiB, 2048px, jpeg+png). We re-encode to JPEG well under all of them.
const (
	// DefaultMaxDim is the longest-edge ceiling. 2048 matches ollama-cloud's
	// MaxDim; anything larger is downscaled. A packet photo has plenty of detail
	// left at 2048.
	DefaultMaxDim = 2048
	// DefaultMaxBytes caps the *input* we will read. A phone photo is 3–8 MB; 25
	// MiB leaves headroom for a large HEIC without inviting a decompression bomb
	// as an unbounded read. The re-encoded output is far smaller.
	DefaultMaxBytes = 25 << 20
	// maxDecodePixels bounds the DECODED bitmap regardless of input byte size, so
	// a small file claiming enormous dimensions (a decompression bomb) is refused
	// before its ~4-bytes/px bitmap is allocated. 50 MP ≈ 200 MB peak — above any
	// current phone camera (a 48 MP sensor is 48 MP) while capping the
	// amplification a hostile header can force. maxDecodePixels and maxDimension
	// are internal safety floors, not knobs — unlike MaxDim/MaxBytes there's no
	// reason for a caller to raise them.
	maxDecodePixels = 50_000_000
	// maxDimension caps EACH side independently. It exists to make the pixel-count
	// check overflow-safe: without it, a header claiming ~2^32 on a side could
	// wrap int64(w)*int64(h) negative and slip past maxDecodePixels. No real image
	// is 50k px on a side.
	maxDimension = 50_000
	// jpegQuality for the normalized output. 85 is visually clean and keeps the
	// file small; the model reads text off it, not fine gradients.
	jpegQuality = 85
)

// Options tunes Normalize. The zero value uses the Default* constants.
type Options struct {
	MaxDim   int // longest edge; 0 → DefaultMaxDim
	MaxBytes int // input read cap; 0 → DefaultMaxBytes
}

func (o Options) maxDim() int {
	if o.MaxDim > 0 {
		return o.MaxDim
	}
	return DefaultMaxDim
}

func (o Options) maxBytes() int {
	if o.MaxBytes > 0 {
		return o.MaxBytes
	}
	return DefaultMaxBytes
}

// ErrTooLarge means the input exceeded the byte cap or decoded to an absurd
// pixel count. ErrUnsupported means the bytes weren't a decodable image format —
// or a decoder panicked on them (see the recover in Normalize).
var (
	ErrTooLarge    = errors.New("imagenorm: image too large")
	ErrUnsupported = errors.New("imagenorm: unsupported or corrupt image")
)

// Normalize reads an image of any supported format (JPEG, PNG, HEIC, WebP),
// downscales it to fit opts.MaxDim on its longest edge, applies the JPEG EXIF
// orientation so the pixels come out upright, and returns it re-encoded as JPEG,
// plus the decoded format name (e.g. "heic") — handy for logging what a phone
// actually sent. On any error the returned bytes are nil and format is "".
//
// Errors, by cause:
//   - input over opts.MaxBytes, or a decoded canvas over maxDecodePixels /
//     maxDimension → ErrTooLarge, refused before the bitmap is allocated;
//   - bytes that aren't a decodable image, or a decoder that panics on them →
//     ErrUnsupported;
//   - a genuine read or JPEG-encode I/O failure → a wrapped error (not a
//     sentinel), since those are the caller's stream/environment, not the image.
//
// It bounds work against a hostile upload three ways: the byte cap, the
// pre-decode pixel/dimension check, and a recover around the third-party decoders
// (a malformed HEIC/WebP shouldn't take the process down).
//
// EXIF orientation: phone cameras store the sensor pixels in one orientation and
// set an EXIF tag to rotate on display, so a JPEG "portrait" photo is really a
// landscape bitmap tagged "rotate 90°" — and the re-encode below strips EXIF,
// which is exactly why the rotation must be BAKED IN here. applyOrientation does
// that for the JPEG path (the format phone uploads overwhelmingly arrive in);
// other formats carry no JPEG EXIF and their decoders own orientation, so they're
// left as decoded.
//
// One known gap, deferred to the upload handler (#81): there is no context —
// image.Decode is CPU-bound and not cancellable mid-decode, so a caller that
// needs a hard deadline should run Normalize under its own timeout. The size
// guards keep the work finite regardless.
func Normalize(r io.Reader, opts Options) (out []byte, format string, err error) {
	// Cap the read at MaxBytes+1 so we can tell "exactly at the cap" from "over".
	// maxBytes() is always a sane positive (default 25 MiB); guard the +1 anyway.
	byteCap := opts.maxBytes()
	limit := int64(byteCap) + 1
	if limit < 1 {
		limit = int64(DefaultMaxBytes) + 1
	}
	raw, err := io.ReadAll(io.LimitReader(r, limit))
	if err != nil {
		return nil, "", fmt.Errorf("imagenorm: read: %w", err)
	}
	if len(raw) > byteCap {
		return nil, "", ErrTooLarge
	}

	// Check dimensions BEFORE a full decode, so a decompression bomb is refused
	// before it allocates its bitmap. The per-side maxDimension check runs first
	// so the pixel-count multiply below can't overflow.
	cfg, _, err := image.DecodeConfig(bytes.NewReader(raw))
	if err != nil {
		return nil, "", ErrUnsupported
	}
	if cfg.Width <= 0 || cfg.Height <= 0 ||
		cfg.Width > maxDimension || cfg.Height > maxDimension ||
		int64(cfg.Width)*int64(cfg.Height) > maxDecodePixels {
		return nil, "", ErrTooLarge
	}

	img, format, err := decodeSafely(raw)
	if err != nil {
		return nil, "", err
	}

	// Downscale first (cheaper to rotate the small image), then bake in the EXIF
	// orientation so the JPEG we emit is upright. A 90° rotation swaps the sides
	// but not the longest edge, so the downscale bound still holds after it.
	img = downscale(img, opts.maxDim())
	img = applyOrientation(img, exifOrientation(raw))

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: jpegQuality}); err != nil {
		return nil, "", fmt.Errorf("imagenorm: encode jpeg: %w", err)
	}
	return buf.Bytes(), format, nil
}

// decodeSafely decodes raw, converting both a decode error and a decoder PANIC
// into ErrUnsupported. The recover matters because the image comes from an
// untrusted upload and the HEIC/WebP decoders are third-party (libheif via WASM,
// x/image/webp): a malformed file that panics one of them must fail this one
// request, not crash the process.
func decodeSafely(raw []byte) (img image.Image, format string, err error) {
	defer func() {
		if r := recover(); r != nil {
			img, format, err = nil, "", ErrUnsupported
		}
	}()
	img, format, err = image.Decode(bytes.NewReader(raw))
	if err != nil {
		return nil, "", ErrUnsupported
	}
	return img, format, nil
}

// applyOrientation returns img with the EXIF orientation (1..8) baked in, so the
// pixels are upright and no display-time rotation is needed. Orientation 1 (and
// anything out of range) is a no-op. Values 5..8 are 90° rotations, which swap
// the output's width and height. Copies raw RGBA pixels by byte offset (after a
// one-time conversion if the source isn't already RGBA), so a full-resolution
// rotation doesn't box a color.Color per pixel.
func applyOrientation(img image.Image, o int) image.Image {
	if o <= 1 || o > 8 {
		return img
	}
	// Work on a concrete RGBA so the transform is a 4-byte copy per pixel rather
	// than millions of boxed color.Color values through At/Set. downscale usually
	// hands us an *image.RGBA already; convert once if not (e.g. a small JPEG that
	// skipped downscale decodes to YCbCr).
	src, ok := img.(*image.RGBA)
	if !ok {
		b := img.Bounds()
		conv := image.NewRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
		draw.Draw(conv, conv.Bounds(), img, b.Min, draw.Src)
		src = conv
	}
	sb := src.Bounds()
	w, h := sb.Dx(), sb.Dy()
	// A quarter-turn (5..8) transposes the output; size the one buffer accordingly.
	dw, dh := w, h
	if o >= 5 {
		dw, dh = h, w
	}
	dst := image.NewRGBA(image.Rect(0, 0, dw, dh))
	for y := range h {
		for x := range w {
			var dx, dy int
			switch o {
			case 2: // mirror horizontal
				dx, dy = w-1-x, y
			case 3: // rotate 180
				dx, dy = w-1-x, h-1-y
			case 4: // mirror vertical
				dx, dy = x, h-1-y
			case 5: // transpose (mirror across the main diagonal)
				dx, dy = y, x
			case 6: // rotate 90° clockwise
				dx, dy = h-1-y, x
			case 7: // transverse (mirror across the anti-diagonal)
				dx, dy = h-1-y, w-1-x
			case 8: // rotate 90° counter-clockwise
				dx, dy = y, w-1-x
			}
			si := src.PixOffset(sb.Min.X+x, sb.Min.Y+y)
			di := dst.PixOffset(dx, dy)
			copy(dst.Pix[di:di+4], src.Pix[si:si+4])
		}
	}
	return dst
}

// exifOrientation extracts the EXIF Orientation tag (1..8) from raw image bytes,
// returning 1 (normal) when it's absent or unparseable — the safe default, since
// a wrong guess rotates a correct image. Only the JPEG APP1/Exif path is parsed:
// that's the format uploaded phone photos overwhelmingly arrive in, and the other
// decoders own their own orientation.
func exifOrientation(raw []byte) int {
	// A JPEG is a run of FFxx marker segments after the SOI (FFD8). Walk them
	// looking for APP1 (FFE1) carrying "Exif\0\0"; stop at the scan data (SOS).
	if len(raw) < 4 || raw[0] != 0xFF || raw[1] != 0xD8 {
		return 1
	}
	for i := 2; i+1 < len(raw); {
		if raw[i] != 0xFF {
			return 1 // not aligned on a marker; give up rather than misread
		}
		// A marker may be preceded by any number of 0xFF fill bytes (JPEG spec);
		// skip them so a padded APP1 isn't misread as a marker of value 0xFF.
		for i+1 < len(raw) && raw[i+1] == 0xFF {
			i++
		}
		if i+1 >= len(raw) {
			return 1
		}
		marker := raw[i+1]
		if marker == 0xD9 || marker == 0xDA {
			return 1 // EOI / start-of-scan: no more headers to read
		}
		if i+4 > len(raw) {
			return 1
		}
		segLen := int(raw[i+2])<<8 | int(raw[i+3])
		if segLen < 2 || i+2+segLen > len(raw) {
			return 1
		}
		if marker == 0xE1 {
			if o, ok := orientationFromApp1(raw[i+4 : i+2+segLen]); ok {
				return o
			}
		}
		i += 2 + segLen
	}
	return 1
}

// orientationFromApp1 reads the Orientation tag from a JPEG APP1 segment body
// (everything after the 2-byte length): "Exif\0\0" then a TIFF block holding
// IFD0. Returns (0, false) if the segment isn't Exif or the tag is missing.
func orientationFromApp1(seg []byte) (int, bool) {
	const prefix = "Exif\x00\x00"
	if len(seg) < len(prefix)+8 || string(seg[:len(prefix)]) != prefix {
		return 0, false
	}
	tiff := seg[len(prefix):]
	var bo binary.ByteOrder
	switch string(tiff[0:2]) {
	case "II":
		bo = binary.LittleEndian
	case "MM":
		bo = binary.BigEndian
	default:
		return 0, false
	}
	if bo.Uint16(tiff[2:4]) != 0x2A { // TIFF magic (42); byte order must agree
		return 0, false
	}
	ifd := int(bo.Uint32(tiff[4:8])) // offset to IFD0 from the TIFF start
	if ifd < 8 || ifd+2 > len(tiff) {
		return 0, false
	}
	n := int(bo.Uint16(tiff[ifd : ifd+2]))
	for k := range n {
		off := ifd + 2 + k*12 // each IFD entry is 12 bytes
		if off+12 > len(tiff) {
			return 0, false
		}
		if bo.Uint16(tiff[off:off+2]) != 0x0112 { // Orientation tag
			continue
		}
		// Orientation is defined as a single SHORT, whose value sits inline in the
		// first 2 bytes of the value field. Reject anything else rather than read a
		// mistyped entry (a LONG/offset there would be a different number entirely).
		if bo.Uint16(tiff[off+2:off+4]) != 3 || bo.Uint32(tiff[off+4:off+8]) != 1 {
			return 0, false
		}
		v := int(bo.Uint16(tiff[off+8 : off+10]))
		if v >= 1 && v <= 8 {
			return v, true
		}
		return 0, false
	}
	return 0, false
}

// downscale returns img shrunk so its longest edge is at most maxDim, preserving
// aspect ratio. An image already within bounds is returned unchanged (no
// re-sampling, no quality loss beyond the JPEG round-trip). Uses Catmull-Rom for
// a sharp result on text, which is what a packet photo is mostly made of.
func downscale(img image.Image, maxDim int) image.Image {
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	longest := max(w, h)
	if longest <= maxDim || longest == 0 {
		return img
	}
	scale := float64(maxDim) / float64(longest)
	nw, nh := max(int(float64(w)*scale), 1), max(int(float64(h)*scale), 1)
	dst := image.NewRGBA(image.Rect(0, 0, nw, nh))
	draw.CatmullRom.Scale(dst, dst.Bounds(), img, b, draw.Over, nil)
	return dst
}