materialize.go

package decodedsource

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sync"
)

var (
	// ErrMaterializedTooLarge reports that decoded output exceeded MaxBytes.
	ErrMaterializedTooLarge = errors.New("materialized output exceeds size limit")
	// ErrClosed reports use of a closed Materialized source.
	ErrClosed = errors.New("materialized source is closed")
)

// MaterializeOptions configures decoded output materialization.
type MaterializeOptions struct {
	// Compression selects the input envelope. The zero value auto-detects it.
	Compression Compression
	// TempDir selects the directory for an automatically managed temporary file.
	// The system temporary directory is used when empty.
	TempDir string
	// Destination creates a persistent file at this path. It takes precedence
	// over TempDir.
	Destination string
	// Mode sets output permissions. The default is 0600.
	Mode fs.FileMode
	// MaxBytes limits decoded output. Zero means unlimited.
	MaxBytes int64
	// Overwrite permits replacing Destination when it already exists.
	Overwrite bool
}

// PersistOptions configures ownership transfer to a persistent path.
type PersistOptions struct {
	// Overwrite permits replacing an existing destination.
	Overwrite bool
}

// Materialized is decoded output stored in a file. It implements a
// random-access source method set without importing a consumer package.
//
// Close removes an automatically managed temporary file. Once persisted,
// Close releases this object but leaves the file owned by the caller.
type Materialized struct {
	mu         sync.RWMutex
	path       string
	persistent bool
	closed     bool
}

// Materialize decodes r once and stores the uncompressed bytes in a seekable
// file. It does not close r.
func Materialize(ctx context.Context, r io.Reader, opts MaterializeOptions) (*Materialized, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if r == nil {
		return nil, errors.New("nil input reader")
	}
	if opts.MaxBytes < 0 {
		return nil, errors.New("MaxBytes must not be negative")
	}

	mode := opts.Mode
	if mode == 0 {
		mode = 0o600
	}
	dir, pattern := opts.TempDir, "decodedsource-*"
	if opts.Destination != "" {
		dir = filepath.Dir(opts.Destination)
		pattern = "." + filepath.Base(opts.Destination) + ".*.tmp"
	}

	file, err := os.CreateTemp(dir, pattern)
	if err != nil {
		return nil, err
	}
	tempPath := file.Name()
	keepTemp := false
	defer func() {
		if !keepTemp {
			_ = os.Remove(tempPath)
		}
	}()

	if err := file.Chmod(mode); err != nil {
		_ = file.Close()
		return nil, err
	}
	decoded, err := NewReader(r, opts.Compression)
	if err != nil {
		_ = file.Close()
		return nil, err
	}
	copyErr := copyWithContext(ctx, file, decoded, opts.MaxBytes)
	decoderCloseErr := decoded.Close()
	if copyErr == nil {
		copyErr = decoderCloseErr
	}
	if copyErr == nil {
		copyErr = file.Sync()
	}
	fileCloseErr := file.Close()
	if copyErr == nil {
		copyErr = fileCloseErr
	}
	if copyErr != nil {
		return nil, copyErr
	}

	materialized := &Materialized{path: tempPath}
	if opts.Destination == "" {
		keepTemp = true
		return materialized, nil
	}
	if err := installFile(tempPath, opts.Destination, opts.Overwrite); err != nil {
		return nil, err
	}
	materialized.path = opts.Destination
	materialized.persistent = true
	return materialized, nil
}

func copyWithContext(ctx context.Context, dst io.Writer, src io.Reader, maxBytes int64) error {
	buffer := make([]byte, 128<<10)
	var written int64
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		n, readErr := src.Read(buffer)
		if n > 0 {
			if maxBytes > 0 && written+int64(n) > maxBytes {
				return fmt.Errorf("%w: limit %d bytes", ErrMaterializedTooLarge, maxBytes)
			}
			wn, writeErr := dst.Write(buffer[:n])
			written += int64(wn)
			if writeErr != nil {
				return writeErr
			}
			if wn != n {
				return io.ErrShortWrite
			}
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return nil
			}
			return readErr
		}
	}
}

// Source returns m for convenient use as a random-access source.
func (m *Materialized) Source() *Materialized { return m }

// Path returns the current filesystem path.
func (m *Materialized) Path() string {
	if m == nil {
		return ""
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.path
}

// Persistent reports whether Close will leave the materialized file intact.
func (m *Materialized) Persistent() bool {
	if m == nil {
		return false
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.persistent
}

// Persist moves the materialized file to path and transfers file ownership to
// the caller. Close no longer removes it. Persist should be called before
// opening scanners or after all readers from Source have been closed.
func (m *Materialized) Persist(path string, opts PersistOptions) error {
	if m == nil {
		return ErrClosed
	}
	if path == "" {
		return errors.New("empty persistence path")
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.closed {
		return ErrClosed
	}
	if filepath.Clean(path) == filepath.Clean(m.path) {
		m.persistent = true
		return nil
	}
	if err := moveFile(m.path, path, opts.Overwrite); err != nil {
		return err
	}
	m.path = path
	m.persistent = true
	return nil
}

// Open opens the complete decoded file.
func (m *Materialized) Open() (io.ReadCloser, error) {
	path, err := m.usablePath()
	if err != nil {
		return nil, err
	}
	return os.Open(path)
}

// OpenRange opens size bytes at off in the decoded file.
func (m *Materialized) OpenRange(off, size int64) (io.ReadCloser, error) {
	if off < 0 || size < 0 {
		return nil, fmt.Errorf("invalid range off=%d size=%d", off, size)
	}
	path, err := m.usablePath()
	if err != nil {
		return nil, err
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	return &sectionReadCloser{Reader: io.NewSectionReader(file, off, size), close: file.Close}, nil
}

// OpenAt opens the decoded file from off through EOF.
func (m *Materialized) OpenAt(off int64) (io.ReadCloser, error) {
	if off < 0 {
		return nil, fmt.Errorf("invalid offset %d", off)
	}
	path, err := m.usablePath()
	if err != nil {
		return nil, err
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	stat, err := file.Stat()
	if err != nil {
		_ = file.Close()
		return nil, err
	}
	if off > stat.Size() {
		off = stat.Size()
	}
	return &sectionReadCloser{Reader: io.NewSectionReader(file, off, stat.Size()-off), close: file.Close}, nil
}

// Size returns the decoded file size.
func (m *Materialized) Size() (int64, error) {
	path, err := m.usablePath()
	if err != nil {
		return 0, err
	}
	stat, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	return stat.Size(), nil
}

func (m *Materialized) usablePath() (string, error) {
	if m == nil {
		return "", ErrClosed
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.closed {
		return "", ErrClosed
	}
	return m.path, nil
}

// Close releases m. Temporary materializations are removed; persisted files
// remain at Path and are owned by the caller.
func (m *Materialized) Close() error {
	if m == nil {
		return nil
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.closed {
		return nil
	}
	m.closed = true
	if m.persistent {
		return nil
	}
	err := os.Remove(m.path)
	if errors.Is(err, fs.ErrNotExist) {
		return nil
	}
	return err
}

type sectionReadCloser struct {
	io.Reader
	close func() error
}

func (r *sectionReadCloser) Close() error {
	if r.close == nil {
		return nil
	}
	return r.close()
}

func installFile(tempPath, destination string, overwrite bool) error {
	if overwrite {
		return os.Rename(tempPath, destination)
	}
	if err := os.Link(tempPath, destination); err != nil {
		return err
	}
	return os.Remove(tempPath)
}

func moveFile(source, destination string, overwrite bool) error {
	if overwrite {
		if err := os.Rename(source, destination); err == nil {
			return nil
		}
	} else {
		if err := os.Link(source, destination); err == nil {
			return os.Remove(source)
		}
	}
	return copyFileAndRemove(source, destination, overwrite)
}

func copyFileAndRemove(source, destination string, overwrite bool) error {
	src, err := os.Open(source)
	if err != nil {
		return err
	}
	stat, err := src.Stat()
	if err != nil {
		_ = src.Close()
		return err
	}

	temp, err := os.CreateTemp(filepath.Dir(destination), "."+filepath.Base(destination)+".*.tmp")
	if err != nil {
		_ = src.Close()
		return err
	}
	tempPath := temp.Name()
	installed := false
	defer func() {
		if !installed {
			_ = os.Remove(tempPath)
		}
	}()
	if err := temp.Chmod(stat.Mode().Perm()); err != nil {
		_ = src.Close()
		_ = temp.Close()
		return err
	}
	if _, err := io.Copy(temp, src); err != nil {
		_ = src.Close()
		_ = temp.Close()
		return err
	}
	if err := src.Close(); err != nil {
		_ = temp.Close()
		return err
	}
	if err := temp.Sync(); err != nil {
		_ = temp.Close()
		return err
	}
	if err := temp.Close(); err != nil {
		return err
	}
	if err := installFile(tempPath, destination, overwrite); err != nil {
		return err
	}
	installed = true
	return os.Remove(source)
}