imagedecoder_heif.go

// Copyright 2026 Toni Melisma
// SPDX-License-Identifier: MIT

package imagemeta

import (
	"io"
	"math"
)

// ISOBMFF box and item types used in HEIF/AVIF containers.
var heifFCC = struct {
	ftyp, meta                         fourCC
	iinf, infe, iloc                   fourCC
	iprp, ipco, ipma, ispe, irot, pitm fourCC
	exif, mime                         fourCC
	moov, trak, mdia, hdlr, minf, stbl fourCC
	stsz, stz2, pict                   fourCC
}{
	ftyp: fourCC{'f', 't', 'y', 'p'},
	meta: fourCC{'m', 'e', 't', 'a'},
	iinf: fourCC{'i', 'i', 'n', 'f'},
	infe: fourCC{'i', 'n', 'f', 'e'},
	iloc: fourCC{'i', 'l', 'o', 'c'},
	iprp: fourCC{'i', 'p', 'r', 'p'},
	ipco: fourCC{'i', 'p', 'c', 'o'},
	ipma: fourCC{'i', 'p', 'm', 'a'},
	ispe: fourCC{'i', 's', 'p', 'e'},
	irot: fourCC{'i', 'r', 'o', 't'},
	pitm: fourCC{'p', 'i', 't', 'm'},
	exif: fourCC{'E', 'x', 'i', 'f'},
	mime: fourCC{'m', 'i', 'm', 'e'},
	moov: fourCC{'m', 'o', 'o', 'v'},
	trak: fourCC{'t', 'r', 'a', 'k'},
	mdia: fourCC{'m', 'd', 'i', 'a'},
	hdlr: fourCC{'h', 'd', 'l', 'r'},
	minf: fourCC{'m', 'i', 'n', 'f'},
	stbl: fourCC{'s', 't', 'b', 'l'},
	stsz: fourCC{'s', 't', 's', 'z'},
	stz2: fourCC{'s', 't', 'z', '2'},
	pict: fourCC{'p', 'i', 'c', 't'},
}

type imageDecoderHEIF struct {
	*baseStreamingDecoder
}

func (e *imageDecoderHEIF) decode() error {
	// readVarUint reads n bytes from the stream as a big-endian uint64.
	// n must be 0, 2, 4, or 8. Returns 0 for n == 0.
	readVarUint := func(n int) uint64 {
		switch n {
		case 0:
			return 0
		case 2:
			return uint64(e.read2())
		case 4:
			return uint64(e.read4())
		case 8:
			return e.read8r(e.r)
		default:
			panic(newInvalidFormatErrorf("heif: unsupported iloc field size: %d", n))
		}
	}

	// readBox reads an ISOBMFF box header from the current stream position.
	// Returns (startPos, totalBoxSize, boxType).
	// startPos: absolute stream position before the header.
	// totalBoxSize: total box size including header bytes (0 = extends to EOF).
	// After this call, the stream is positioned at the start of the box payload.
	readBox := func() (startPos int64, totalSize uint64, boxType fourCC) {
		startPos = e.pos()
		size := e.read4()
		e.readBytes(boxType[:])
		totalSize = uint64(size)
		if size == 1 {
			// Extended size: next 8 bytes hold the actual size.
			totalSize = e.read8r(e.r)
		}
		return
	}

	// Step 1: Read and validate the ftyp box.
	ftypStart, ftypSize, ftypType := readBox()
	if e.isEOF {
		return errInvalidFormat
	}
	if ftypType != heifFCC.ftyp {
		return errInvalidFormat
	}
	if ftypSize > 0 {
		e.seek(ftypStart + int64(ftypSize))
	}

	// eachBox calls f for each box up to end, with the stream positioned at the box payload.
	// It stops at the first box with a size not covering its header or exceeding end.
	eachBox := func(end int64, f func(boxType fourCC, boxEnd int64)) {
		for e.pos()+8 <= end {
			start, size, boxType := readBox()
			if e.isEOF || int64(size) < e.pos()-start || size > uint64(end-start) {
				return
			}
			boxEnd := start + int64(size)
			f(boxType, boxEnd)
			e.seek(boxEnd)
		}
	}

	// Step 2: Scan top-level boxes for the meta box and,
	// for CONFIG, the moov box of image sequences (e.g. animated AVIF).
	needMoov := e.opts.Sources.Has(CONFIG)
	var (
		metaStart   int64
		metaSize    uint64
		foundMeta   bool
		moovPayload int64
		moovEnd     int64
		foundMoov   bool
	)
	pos := e.pos()
	streamEnd, err := e.r.Seek(0, io.SeekEnd)
	if err != nil {
		return err
	}
	e.seek(pos)
	for e.pos()+8 <= streamEnd {
		s, size, boxType := readBox()
		if size != 0 && int64(size) < e.pos()-s {
			break
		}
		switch boxType {
		case heifFCC.meta:
			metaStart, metaSize, foundMeta = s, size, true
		case heifFCC.moov:
			moovPayload, moovEnd, foundMoov = e.pos(), streamEnd, true
			if size != 0 {
				moovEnd = min(moovEnd, s+int64(size))
			}
		}
		if size == 0 || (foundMeta && (foundMoov || !needMoov)) {
			break
		}
		e.seek(s + int64(size))
	}

	if foundMoov && needMoov {
		// The frame count is the sample count of the image sequence ('pict') track.
		e.seek(moovPayload)
		eachBox(moovEnd, func(boxType fourCC, boxEnd int64) {
			if boxType != heifFCC.trak {
				return
			}
			var (
				handler fourCC
				count   uint32
			)
			eachBox(boxEnd, func(boxType fourCC, boxEnd int64) {
				if boxType != heifFCC.mdia {
					return
				}
				eachBox(boxEnd, func(boxType fourCC, boxEnd int64) {
					switch boxType {
					case heifFCC.hdlr:
						e.skip(8) // version+flags, pre_defined
						e.readBytes(handler[:])
					case heifFCC.minf:
						eachBox(boxEnd, func(boxType fourCC, boxEnd int64) {
							if boxType != heifFCC.stbl {
								return
							}
							eachBox(boxEnd, func(boxType fourCC, boxEnd int64) {
								if boxType == heifFCC.stsz || boxType == heifFCC.stz2 {
									e.skip(8) // version+flags, sample_size (stsz) or field_size (stz2)
									count = e.read4()
								}
							})
						})
					}
				})
			})
			if handler == heifFCC.pict {
				e.result.ImageConfig.FrameCount = max(e.result.ImageConfig.FrameCount, int(count))
			}
		})
	}

	if !foundMeta {
		return nil
	}
	e.seek(metaStart)
	readBox()

	// Step 3: Parse the meta FullBox (skip 4 bytes version+flags).
	e.skip(4)

	var metaEnd int64
	if metaSize == 0 {
		metaEnd = math.MaxInt64 // extends to EOF
	} else {
		metaEnd = metaStart + int64(metaSize)
	}

	var (
		exifItemID    uint32
		xmpItemIDs    []uint32
		primaryItemID uint32
	)

	// iloc entries keyed by item ID, resolved after the full meta scan
	// so that box ordering (iloc before/after iinf) doesn't matter.
	type ilocEntry struct {
		offset, length uint64
	}
	needEXIFOrXMP := e.opts.Sources.Has(EXIF) || e.opts.Sources.Has(XMP)
	var ilocEntries map[uint32]ilocEntry

	// For CONFIG: ipco properties and ipma associations are collected during
	// the meta box scan and resolved afterwards, so box ordering doesn't matter.
	type ipcoProp struct {
		isIspe        bool
		isIrot        bool
		width, height uint32
		angle         uint8
	}
	var ipcoProps []ipcoProp
	var primaryPropIndices []int // 1-based property indices from ipma

	// Step 4: Iterate inner boxes of meta to find pitm, iinf, iloc, and iprp.
	for e.pos()+8 <= metaEnd {
		innerStart, innerSize, innerType := readBox()
		if e.isEOF {
			break
		}
		if innerSize == 0 {
			break // extends to EOF
		}
		innerEnd := innerStart + int64(innerSize)

		switch innerType {
		case heifFCC.pitm:
			if e.opts.Sources.Has(CONFIG) {
				vf := e.read4()
				if vf>>24 == 0 {
					primaryItemID = uint32(e.read2())
				} else {
					primaryItemID = e.read4()
				}
			}

		case heifFCC.iinf:
			if !needEXIFOrXMP {
				break
			}
			// iinf is a FullBox: read version+flags then item count.
			vf := e.read4()
			iinfVersion := vf >> 24
			var count uint32
			if iinfVersion == 0 {
				count = uint32(e.read2())
			} else {
				count = e.read4()
			}

			// Iterate infe sub-boxes.
			for range count {
				infeStart, infeSize, infeType := readBox()
				if e.isEOF || infeSize == 0 {
					break
				}
				infeEnd := infeStart + int64(infeSize)

				if infeType == heifFCC.infe {
					// infe is a FullBox: read version+flags.
					vf2 := e.read4()
					infeVersion := vf2 >> 24
					if infeVersion >= 2 {
						var itemID uint32
						if infeVersion == 2 {
							itemID = uint32(e.read2())
						} else {
							// Version 3: 32-bit item ID.
							itemID = e.read4()
						}
						e.skip(2) // protectionIndex
						var itemType fourCC
						e.readBytes(itemType[:])
						switch itemType {
						case heifFCC.exif:
							exifItemID = itemID
						case heifFCC.mime:
							// A "mime" infe carries item_name and content_type as
							// null-terminated strings. XMP is content_type
							// "application/rdf+xml"; other mime items (ICC profiles,
							// etc.) must be ignored.
							remaining := int(infeEnd - e.pos())
							if remaining > 0 {
								_, n := e.readNullTerminatedBytes(remaining)
								remaining -= int(n)
							}
							if remaining > 0 {
								contentType, _ := e.readNullTerminatedBytes(remaining)
								if string(contentType) == "application/rdf+xml" {
									xmpItemIDs = append(xmpItemIDs, itemID)
								}
							}
						}
					} else {
						e.opts.Warnf("heif: infe version %d not supported, skipping", infeVersion)
					}
				}
				e.seek(infeEnd)
			}

		case heifFCC.iloc:
			if !needEXIFOrXMP {
				break
			}
			if ilocEntries == nil {
				ilocEntries = make(map[uint32]ilocEntry)
			}
			// iloc is a FullBox: read version+flags.
			vf := e.read4()
			ilocVersion := uint8(vf >> 24)

			b1 := e.read1()
			offsetSize := int(b1 >> 4)
			lengthSize := int(b1 & 0x0f)

			b2 := e.read1()
			baseOffsetSize := int(b2 >> 4)
			indexSize := int(b2 & 0x0f)

			var count uint32
			if ilocVersion < 2 {
				count = uint32(e.read2())
			} else {
				count = e.read4()
			}

			for range count {
				var itemID uint32
				if ilocVersion < 2 {
					itemID = uint32(e.read2())
				} else {
					itemID = e.read4()
				}

				var constructionMethod uint16
				if ilocVersion >= 1 {
					constructionMethod = e.read2()
				}
				e.skip(2) // dataReferenceIndex

				baseOffset := readVarUint(baseOffsetSize)

				extentCount := e.read2()

				// Only file-offset construction (method 0) is supported.
				if constructionMethod != 0 {
					for range extentCount {
						if ilocVersion >= 1 && indexSize > 0 {
							readVarUint(indexSize)
						}
						readVarUint(offsetSize)
						readVarUint(lengthSize)
					}
					continue
				}

				var firstOffset, firstLength uint64
				for j := range extentCount {
					if ilocVersion >= 1 && indexSize > 0 {
						readVarUint(indexSize) // extent index, discard
					}
					off := readVarUint(offsetSize)
					length := readVarUint(lengthSize)
					if j == 0 {
						firstOffset = baseOffset + off
						firstLength = length
					}
				}

				ilocEntries[itemID] = ilocEntry{offset: firstOffset, length: firstLength}
			}

		case heifFCC.iprp:
			if e.opts.Sources.Has(CONFIG) {
				iprpEnd := innerEnd
				for e.pos()+8 <= iprpEnd {
					childStart, childSize, childType := readBox()
					if e.isEOF || childSize == 0 {
						break
					}
					childEnd := childStart + int64(childSize)

					switch childType {
					case heifFCC.ipco:
						for e.pos()+8 <= childEnd {
							propStart, propSize, propType := readBox()
							if e.isEOF || propSize == 0 {
								break
							}
							propEnd := propStart + int64(propSize)

							var prop ipcoProp
							switch propType {
							case heifFCC.ispe:
								e.skip(4) // version+flags
								prop = ipcoProp{isIspe: true, width: e.read4(), height: e.read4()}
							case heifFCC.irot:
								prop = ipcoProp{isIrot: true, angle: e.read1()}
							}
							ipcoProps = append(ipcoProps, prop)
							e.seek(propEnd)
						}

					case heifFCC.ipma:
						// ipma maps item IDs to property indices in ipco.
						vf := e.read4()
						ipmaVersion := uint8(vf >> 24)
						ipmaFlags := vf & 0xFFFFFF
						entryCount := e.read4()
						for range entryCount {
							var itemID uint32
							if ipmaVersion < 1 {
								itemID = uint32(e.read2())
							} else {
								itemID = e.read4()
							}
							assocCount := e.read1()
							for range assocCount {
								var propIdx int
								if ipmaFlags&1 != 0 {
									val := e.read2()
									propIdx = int(val & 0x7FFF)
								} else {
									val := e.read1()
									propIdx = int(val & 0x7F)
								}
								if itemID == primaryItemID && primaryItemID != 0 {
									primaryPropIndices = append(primaryPropIndices, propIdx)
								}
							}
						}
					}
					e.seek(childEnd)
				}
			}
		}

		// Always advance to the end of this inner box.
		e.seek(innerEnd)
	}

	// Step 5: Resolve iloc offsets now that both iinf and iloc have been parsed.
	var exifOffset, exifLength uint64
	if loc, ok := ilocEntries[exifItemID]; ok && exifItemID != 0 {
		exifOffset, exifLength = loc.offset, loc.length
	}

	// Step 6: Resolve CONFIG dimensions from collected ipco/ipma/pitm data.
	if e.opts.Sources.Has(CONFIG) && len(ipcoProps) > 0 {
		var cfgWidth, cfgHeight uint32
		var cfgRotate bool

		if primaryItemID != 0 && len(primaryPropIndices) > 0 {
			// Primary path: use pitm + ipma to find the primary item's properties.
			for _, idx := range primaryPropIndices {
				if idx < 1 || idx > len(ipcoProps) {
					continue
				}
				p := ipcoProps[idx-1]
				if p.isIspe && p.width > 0 && p.height > 0 {
					cfgWidth, cfgHeight = p.width, p.height
				}
				if p.isIrot && (p.angle == 1 || p.angle == 3) {
					cfgRotate = true
				}
			}
		}

		if cfgWidth == 0 || cfgHeight == 0 {
			// Fallback: use the largest ispe (primary image is always larger
			// than tiles or thumbnails in standard HEIF/AVIF output).
			for _, p := range ipcoProps {
				if p.isIspe && p.width > 0 && p.height > 0 {
					if uint64(p.width)*uint64(p.height) > uint64(cfgWidth)*uint64(cfgHeight) {
						cfgWidth, cfgHeight = p.width, p.height
					}
				}
			}
			for _, p := range ipcoProps {
				if p.isIrot && (p.angle == 1 || p.angle == 3) {
					cfgRotate = true
					break
				}
			}
		}

		if cfgWidth > 0 && cfgHeight > 0 {
			if cfgRotate {
				cfgWidth, cfgHeight = cfgHeight, cfgWidth
			}
			e.result.ImageConfig.Width, e.result.ImageConfig.Height = int(cfgWidth), int(cfgHeight)
		}
	}

	// Step 7: Extract EXIF metadata using the absolute offset from iloc.
	if e.opts.Sources.Has(EXIF) && exifItemID != 0 && exifOffset != 0 && exifLength > 4 {
		if err := e.handleEXIF(exifOffset, exifLength); err != nil {
			return err
		}
	}

	// Step 8: Extract XMP metadata. HEIF permits multiple XMP items (e.g. Ultra
	// HDR: one on the primary image, one on the gain map). Deliver each in file
	// order; HandleXMP is invoked once per packet.
	if e.opts.Sources.Has(XMP) {
		for _, id := range xmpItemIDs {
			loc, ok := ilocEntries[id]
			if !ok || loc.offset == 0 || loc.length == 0 {
				continue
			}
			if err := func() error {
				e.seek(int64(loc.offset))
				r, err := e.bufferedReader(int64(loc.length))
				if err != nil {
					return err
				}
				defer r.Close()
				return decodeXMP(r, e.opts)
			}(); err != nil {
				return err
			}
		}
	}

	return nil
}

func (e *imageDecoderHEIF) handleEXIF(offset, length uint64) (err error) {
	defer func() {
		if r := recover(); r != nil {
			// Recover from panic in EXIF decoder (e.g., errStop).
			// This allows the HEIF decoder to continue processing XMP.
			if rerr, ok := r.(error); ok && rerr != errStop {
				err = rerr
			}
		}
	}()

	e.seek(int64(offset))
	// HEIF EXIF blobs are prefixed with a 4-byte big-endian header offset.
	// Skip that many bytes to reach the TIFF header.
	exifHdrOffset := e.read4()
	if int64(exifHdrOffset) > int64(length)-4 {
		return newInvalidFormatErrorf("heif: invalid exif header offset %d", exifHdrOffset)
	}
	e.skip(int64(exifHdrOffset))
	thumbnailPos := e.pos()
	dataLen := int64(length) - 4 - int64(exifHdrOffset)
	if dataLen <= 0 {
		return nil
	}
	r, err := e.bufferedReader(dataLen)
	if err != nil {
		return err
	}
	defer r.Close()
	dec := newMetaDecoderEXIF(r, e.byteOrder, thumbnailPos, e.opts)
	return dec.decode()
}