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"
	"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, 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).
//
// Two known gaps, both deferred to the upload handler that wires this in (#81):
//   - EXIF ORIENTATION is not applied, so a portrait phone photo tagged
//     "rotate 90°" comes out sideways. That's best fixed and tested with a real
//     oriented photo end-to-end, which the library has no consumer for yet.
//   - 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
	}

	img = downscale(img, opts.maxDim())

	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
}

// 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
}