io.go

// Copyright 2024 Bjørn Erik Pedersen
// SPDX-License-Identifier: MIT

package imagemeta

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"sync"
)

type bufferedReadSeeker struct {
	rs  io.ReadSeeker
	br  *bufio.Reader
	pos int64
}

var bufferedReadSeekerPool = &sync.Pool{
	New: func() any {
		return &bufferedReadSeeker{br: bufio.NewReaderSize(nil, 4096)}
	},
}

func getBufferedReadSeeker(rs io.ReadSeeker) *bufferedReadSeeker {
	b := bufferedReadSeekerPool.Get().(*bufferedReadSeeker)
	b.rs = rs
	b.br.Reset(rs)
	b.pos, _ = rs.Seek(0, io.SeekCurrent)
	return b
}

func putBufferedReadSeeker(b *bufferedReadSeeker) {
	b.rs = nil
	b.br.Reset(nil)
	bufferedReadSeekerPool.Put(b)
}

func (b *bufferedReadSeeker) Read(p []byte) (int, error) {
	n, err := b.br.Read(p)
	b.pos += int64(n)
	return n, err
}

// Seek discards from the buffer when possible, avoiding a reset of the buffer
// when e.g. skipping many small chunks.
func (b *bufferedReadSeeker) Seek(offset int64, whence int) (int64, error) {
	target := offset
	switch whence {
	case io.SeekCurrent:
		target += b.pos
	case io.SeekEnd:
		n, err := b.rs.Seek(offset, whence)
		if err != nil {
			return n, err
		}
		b.br.Reset(b.rs)
		b.pos = n
		return n, nil
	}

	if d := target - b.pos; d >= 0 && d <= int64(b.br.Buffered()) {
		b.br.Discard(int(d))
		b.pos = target
		return target, nil
	}

	n, err := b.rs.Seek(target, io.SeekStart)
	if err != nil {
		return n, err
	}
	b.br.Reset(b.rs)
	b.pos = n
	return n, nil
}

type bytesAndReader struct {
	b []byte
	r *bytes.Reader
}

var bytesAndReaderPool = &sync.Pool{
	New: func() any {
		return &bytesAndReader{
			b: make([]byte, 1024),
			r: bytes.NewReader(nil),
		}
	},
}

func getBytesAndReader(length int) *bytesAndReader {
	b := bytesAndReaderPool.Get().(*bytesAndReader)
	if length > cap(b.b) {
		b.b = make([]byte, length)
	}
	b.b = b.b[:length]
	return b
}

func putBytesAndReader(br *bytesAndReader) {
	br.b = br.b[:0]
	bytesAndReaderPool.Put(br)
}

func (br *bytesAndReader) Read(p []byte) (int, error) {
	return br.r.Read(p)
}

func (br *bytesAndReader) Seek(offset int64, whence int) (int64, error) {
	return br.r.Seek(offset, whence)
}

func (br *bytesAndReader) Close() error {
	putBytesAndReader(br)
	return nil
}

var errShortRead = errors.New("short read")

func newStreamReader(r io.Reader, byteOrder binary.ByteOrder) *streamReader {
	return &streamReader{
		r:         r.(io.ReadSeeker),
		byteOrder: byteOrder,
	}
}

type decoder interface {
	decode() error
}

type fourCC [4]byte

type readerCloser interface {
	io.ReadSeeker
	io.Closer
}

// streamReader is a wrapper around a Reader that provides methods to read binary data.
// Note that this is not thread safe.
type streamReader struct {
	// The current Reader.
	r         io.ReadSeeker
	byteOrder binary.ByteOrder

	buf []byte

	isEOF        bool
	readErr      error
	readerOffset int64
}

func (e *streamReader) otherByteOrder() binary.ByteOrder {
	if e.byteOrder == binary.BigEndian {
		return binary.LittleEndian
	}
	return binary.BigEndian
}

// 10 MB should be plenty for image metadata.
const maxBufSize = 10 * 1024 * 1024

// bufferedReader reads length bytes from the stream and returns a ReaderCloser.
// It's important to call Close on the ReaderCloser when done.
func (e *streamReader) bufferedReader(length int64) (readerCloser, error) {
	if length > maxBufSize {
		return nil, newInvalidFormatErrorf("length %d exceeds max %d", length, maxBufSize)
	}
	if length < 0 {
		return nil, newInvalidFormatErrorf("negative length")
	}

	br := getBytesAndReader(int(length))

	if _, err := io.ReadFull(e.r, br.b); err != nil {
		putBytesAndReader(br)
		return nil, err
	}

	br.r.Reset(br.b)

	return br, nil
}

func (e *streamReader) allocateBuf(length int) {
	if length > cap(e.buf) {
		e.buf = make([]byte, length)
	}
}

func (e *streamReader) pos() int64 {
	n, _ := e.r.Seek(0, 1)
	return n
}

func (e *streamReader) read1() uint8 {
	return e.read1r(e.r)
}

func (e *streamReader) read1r(r io.Reader) uint8 {
	const n = 1
	e.readNFromRIntoBuf(n, r)
	return e.buf[0]
}

func (e *streamReader) read2() uint16 {
	return e.read2r(e.r)
}

func (e *streamReader) read2E() (uint16, error) {
	const n = 2
	if err := e.readNIntoBufE(n); err != nil {
		return 0, err
	}
	return e.byteOrder.Uint16(e.buf[:n]), nil
}

func (e *streamReader) read2r(r io.Reader) uint16 {
	const n = 2
	e.readNFromRIntoBuf(n, r)
	return e.byteOrder.Uint16(e.buf[:n])
}

func (e *streamReader) read4() uint32 {
	const n = 4
	e.readNIntoBuf(n)
	return e.byteOrder.Uint32(e.buf[:n])
}

func (e *streamReader) read4r(r io.Reader) uint32 {
	const n = 4
	e.readNFromRIntoBuf(n, r)
	return e.byteOrder.Uint32(e.buf[:n])
}

func (e *streamReader) read4sr(r io.Reader) int32 {
	const n = 4
	e.readNFromRIntoBuf(n, r)
	return int32(e.byteOrder.Uint32(e.buf[:n]))
}

func (e *streamReader) read8r(r io.Reader) uint64 {
	const n = 8
	e.readNFromRIntoBuf(n, r)
	return e.byteOrder.Uint64(e.buf[:n])
}

func (e *streamReader) readBytes(b []byte) error {
	if _, err := io.ReadFull(e.r, b); err != nil {
		e.stop(err)
	}
	return nil
}

// readNullTerminatedBytes reads a slice of bytes from the stream
// until a null byte is encountered.
// It returns the slice of bytes read and the number of bytes read.
// Note that max is the maximum number of bytes to read including the null byte.
func (e *streamReader) readNullTerminatedBytes(max int) ([]byte, int64) {
	var b []byte
	var n int64
	for i := range max {
		b = append(b, e.read1())
		n++
		if b[i] == 0 {
			return b[:i], n
		}
	}
	return b, n
}

// readBytesVolatile reads a slice of bytes from the stream
// which is not guaranteed to be valid after the next read.
func (e *streamReader) readBytesVolatile(n int) []byte {
	e.readNIntoBuf(n)
	return e.buf[:n]
}

func (e *streamReader) readBytesVolatileE(n int) ([]byte, error) {
	err := e.readNIntoBufE(n)
	if err != nil {
		return nil, err
	}
	return e.buf[:n], nil
}

func (e *streamReader) readBytesFromRVolatile(n int, r io.Reader) []byte {
	e.readNFromRIntoBuf(n, r)
	return e.buf[:n]
}

func (e *streamReader) readNFromRIntoBuf(n int, r io.Reader) {
	if err := e.readNFromRIntoBufE(n, r); err != nil {
		e.stop(err)
	}
}

func (e *streamReader) readNFromRIntoBufE(n int, r io.Reader) error {
	if n < 0 {
		return newInvalidFormatErrorf("negative length %d", n)
	}
	if n > maxBufSize {
		return newInvalidFormatErrorf("length %d exceeds max %d", n, maxBufSize)
	}
	e.allocateBuf(n)
	n2, err := io.ReadFull(r, e.buf[:n])
	if err != nil {
		return err
	}
	if n != n2 {
		return errShortRead
	}
	return nil
}

func (e *streamReader) readNIntoBuf(n int) {
	e.readNFromRIntoBuf(n, e.r)
}

func (e *streamReader) readNIntoBufE(n int) error {
	return e.readNFromRIntoBufE(n, e.r)
}

func (e *streamReader) preservePos(f func() error) error {
	pos := e.pos()
	err := f()
	e.seek(pos)
	return err
}

func (e *streamReader) seek(pos int64) {
	_, err := e.r.Seek(pos, io.SeekStart)
	if err != nil {
		e.stop(err)
	}
}

func (e *streamReader) skip(n int64) {
	e.r.Seek(n, io.SeekCurrent)
}

func (e *streamReader) stop(err error) {
	// Alow one silent EOF.
	// This allows the client to not having to check for EOF on every read.
	if err == io.EOF && !e.isEOF {
		e.isEOF = true
		return
	}
	if err != nil {
		e.readErr = err
	}
	panic(errStop)
}