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 TestNormalize exercises all four formats to keep them so.
package imagenorm

import (
	"bytes"
	"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 — a
	// small file can claim enormous dimensions (a decompression bomb). 100 MP is
	// ~400 MB at 4 bytes/px decoded, already generous for any real photo.
	maxDecodePixels = 100_000_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.
var (
	ErrTooLarge    = fmt.Errorf("imagenorm: image too large")
	ErrUnsupported = fmt.Errorf("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. It also reports the decoded format name (e.g. "heic"), which is handy
// for logging what a phone actually sent.
//
// It reads at most opts.MaxBytes and refuses a bitmap over maxDecodePixels, so a
// hostile upload can neither exhaust memory by byte count nor by claimed
// dimensions. A read/format error maps to ErrUnsupported; an over-limit input to
// ErrTooLarge.
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".
	limited := io.LimitReader(r, int64(opts.maxBytes())+1)
	raw, err := io.ReadAll(limited)
	if err != nil {
		return nil, "", fmt.Errorf("imagenorm: read: %w", err)
	}
	if len(raw) > opts.maxBytes() {
		return nil, "", ErrTooLarge
	}

	// Check dimensions BEFORE a full decode, so a decompression bomb is refused
	// before it allocates its bitmap.
	cfg, format, err := image.DecodeConfig(bytes.NewReader(raw))
	if err != nil {
		return nil, "", ErrUnsupported
	}
	if int64(cfg.Width)*int64(cfg.Height) > maxDecodePixels {
		return nil, format, ErrTooLarge
	}

	img, _, err := image.Decode(bytes.NewReader(raw))
	if err != nil {
		return nil, format, ErrUnsupported
	}

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

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: jpegQuality}); err != nil {
		return nil, format, fmt.Errorf("imagenorm: encode jpeg: %w", err)
	}
	return buf.Bytes(), 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 := w
	if h > longest {
		longest = 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
}