merge.go

package mergewarc

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"maps"
	"os"

	"github.com/saveweb/unwarc"
)

const ManifestVersion = 1

const zstdDictionaryMagic = uint32(0x184d2a5d)

var ErrDictionary = errors.New("mergewarc: zstd dictionaries are not supported")

// Input describes one WARC-Zstd object.
type Input struct {
	Name     string
	Open     func() (io.ReadCloser, error)
	Metadata map[string]string
}

// FileInput constructs an Input backed by a local file.
func FileInput(name, path string) Input {
	return Input{
		Name: name,
		Open: func() (io.ReadCloser, error) {
			return os.Open(path)
		},
	}
}

// Entry identifies one original input inside the merged WARC-Zstd object.
type Entry struct {
	Name      string            `json:"name"`
	Offset    int64             `json:"offset"`
	Size      int64             `json:"size"`
	Checksums []string          `json:"checksums"`
	Records   int64             `json:"records"`
	Metadata  map[string]string `json:"metadata,omitempty"`
}

// Output describes the complete merged object.
type Output struct {
	Size      int64    `json:"size"`
	Checksums []string `json:"checksums"`
}

// Manifest maps source objects to byte ranges in a merged WARC-Zstd object.
type Manifest struct {
	Version int     `json:"version"`
	Format  string  `json:"format"`
	Entries []Entry `json:"entries"`
	Output  Output  `json:"output"`
}

// Merge validates inputs and copies their encoded bytes to dst in order.
// Dictionary-bearing inputs are rejected. If Merge returns an error, dst may
// contain a partial result; callers publishing files should write atomically.
func Merge(ctx context.Context, dst io.Writer, inputs []Input, opts Options) (*Manifest, error) {
	if len(inputs) == 0 {
		return nil, fmt.Errorf("mergewarc: no inputs")
	}
	if dst == nil {
		return nil, fmt.Errorf("mergewarc: nil destination")
	}

	seen := make(map[string]struct{}, len(inputs))
	for i, input := range inputs {
		if input.Name == "" {
			return nil, fmt.Errorf("mergewarc: input %d has an empty name", i)
		}
		if input.Open == nil {
			return nil, fmt.Errorf("mergewarc: input %q has no opener", input.Name)
		}
		if _, ok := seen[input.Name]; ok {
			return nil, fmt.Errorf("mergewarc: duplicate input name %q", input.Name)
		}
		seen[input.Name] = struct{}{}
	}

	outputHashes, err := opts.checksumHashes()
	if err != nil {
		return nil, err
	}
	outputWriters := append([]io.Writer{dst}, hashWriters(outputHashes)...)
	outputWriter := io.MultiWriter(outputWriters...)
	algorithms := make([]ChecksumAlgorithm, 0, len(outputHashes))
	for _, item := range outputHashes {
		algorithms = append(algorithms, item.algorithm)
	}
	manifest := &Manifest{
		Version: ManifestVersion,
		Format:  "warc-zstd",
		Entries: make([]Entry, 0, len(inputs)),
	}

	var offset int64
	for _, input := range inputs {
		entry, err := mergeInput(ctx, outputWriter, input, algorithms)
		if err != nil {
			return nil, fmt.Errorf("merge %q: %w", input.Name, err)
		}
		entry.Offset = offset
		entry.Metadata = cloneMetadata(input.Metadata)
		manifest.Entries = append(manifest.Entries, entry)
		offset += entry.Size
	}

	manifest.Output = Output{
		Size:      offset,
		Checksums: formatChecksums(outputHashes),
	}
	return manifest, nil
}

func mergeInput(ctx context.Context, dst io.Writer, input Input, algorithms []ChecksumAlgorithm) (Entry, error) {
	r, err := input.Open()
	if err != nil {
		return Entry{}, err
	}
	defer r.Close()

	hashes, err := (Options{Checksums: algorithms}).checksumHashes()
	if err != nil {
		return Entry{}, err
	}
	writers := append([]io.Writer{dst}, hashWriters(hashes)...)
	counter := &countingWriter{w: io.MultiWriter(writers...)}
	buffered := bufio.NewReader(ctxReader{ctx: ctx, r: r})
	if prefix, err := buffered.Peek(4); err == nil && binary.LittleEndian.Uint32(prefix) == zstdDictionaryMagic {
		return Entry{}, ErrDictionary
	}
	tee := io.TeeReader(buffered, counter)
	opts := unwarc.DefaultScannerOptions()
	opts.Compression = unwarc.CompressionZstd
	scanner, err := unwarc.NewScanner(tee, opts)
	if err != nil {
		return Entry{}, err
	}
	defer scanner.Close()

	var records int64
	for scanner.Next() {
		records++
	}
	if err := scanner.Err(); err != nil {
		return Entry{}, err
	}
	if records == 0 {
		return Entry{}, fmt.Errorf("contains no WARC records")
	}
	return Entry{
		Name:      input.Name,
		Size:      counter.n,
		Checksums: formatChecksums(hashes),
		Records:   records,
	}, nil
}

// Extract copies and verifies one manifest entry from a merged object.
func Extract(ctx context.Context, src io.ReaderAt, entry Entry, dst io.Writer) error {
	if src == nil || dst == nil {
		return fmt.Errorf("mergewarc: nil extraction source or destination")
	}
	if entry.Offset < 0 || entry.Size < 0 {
		return fmt.Errorf("mergewarc: invalid range for %q", entry.Name)
	}
	hashes, err := parseChecksums(entry.Checksums)
	if err != nil {
		return fmt.Errorf("mergewarc: invalid checksums for %q: %w", entry.Name, err)
	}
	writers := append([]io.Writer{dst}, hashWriters(hashes)...)
	n, err := io.Copy(io.MultiWriter(writers...), ctxReader{
		ctx: ctx,
		r:   io.NewSectionReader(src, entry.Offset, entry.Size),
	})
	if err != nil {
		return err
	}
	if n != entry.Size {
		return fmt.Errorf("mergewarc: extracted %d bytes for %q, expected %d", n, entry.Name, entry.Size)
	}
	for _, item := range hashes {
		got := item.hash.Sum(nil)
		if !equalDigest(got, item.want) {
			return fmt.Errorf("mergewarc: checksum mismatch for %q: got %s, expected %s", entry.Name,
				formatDigest(item.algorithm, got), formatDigest(item.algorithm, item.want))
		}
	}
	return nil
}

func cloneMetadata(src map[string]string) map[string]string {
	if len(src) == 0 {
		return nil
	}
	dst := make(map[string]string, len(src))
	maps.Copy(dst, src)
	return dst
}

type ctxReader struct {
	ctx context.Context
	r   io.Reader
}

func (r ctxReader) Read(p []byte) (int, error) {
	if err := r.ctx.Err(); err != nil {
		return 0, err
	}
	return r.r.Read(p)
}

type countingWriter struct {
	w io.Writer
	n int64
}

func (w *countingWriter) Write(p []byte) (int, error) {
	n, err := w.w.Write(p)
	w.n += int64(n)
	return n, err
}