writer.go

package cpio

import (
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"

	"github.com/aibor/cpio/internal"
)

var (
	// ErrWriteTooLong indicates that an attempt was made to write more than
	// Header.Size bytes to the current file.
	ErrWriteTooLong = errors.New("cpio: write too long")

	// ErrWriteAfterClose indicates that an attempt was made to write to the
	// CPIO archive after it was closed.
	ErrWriteAfterClose = errors.New("cpio: write after close")
)

var trailer = &Header{
	Name:  headerEOF,
	Links: 1,
}

var zeroBlock [4]byte

// Writer provides sequential writing of a CPIO archive. Write.WriteHeader
// begins a new file with the provided Header, and then Writer can be treated as
// an io.Writer to supply that file's data.
type Writer struct {
	w      io.Writer
	nb     int64 // number of unwritten bytes for current file entry
	pad    int64 // amount of padding to write after current file entry
	inode  int64
	err    error
	closed bool
}

// NewWriter creates a new Writer writing to w.
func NewWriter(w io.Writer) *Writer {
	return &Writer{w: w}
}

// Flush finishes writing the current file's block padding. The current file
// must be fully written before Flush can be called.
//
// This is unnecessary as the next call to WriteHeader or Close will implicitly
// flush out the file's padding.
func (w *Writer) Flush() error {
	if w.nb > 0 {
		w.err = fmt.Errorf("cpio: missed writing %d bytes", w.nb)
		return w.err
	}
	_, w.err = w.w.Write(zeroBlock[:w.pad])
	if w.err != nil {
		return w.err
	}
	w.nb = 0
	w.pad = 0
	return w.err
}

// WriteHeader writes hdr and prepares to accept the file's contents. The
// Header.Size determines how many bytes can be written for the next file. If
// the current file is not fully written, then this returns an error. This
// implicitly flushes any padding necessary before writing the header.
func (w *Writer) WriteHeader(hdr *Header) (err error) {
	if w.closed {
		return ErrWriteAfterClose
	}
	if w.err == nil {
		w.Flush()
	}
	if w.err != nil {
		return w.err
	}
	if hdr.Name != headerEOF {
		// TODO: should we be mutating hdr here?
		// ensure all inodes are unique
		w.inode++
		if hdr.Inode == 0 {
			hdr.Inode = w.inode
		}

		// ensure file type is set
		if hdr.Mode&^ModePerm == 0 {
			hdr.Mode |= ModeRegular
		}

		// ensure regular files have at least 1 inbound link
		if hdr.Links < 1 && hdr.Mode.IsRegular() {
			hdr.Links = 1
		}
	}

	w.nb = hdr.Size
	w.pad, w.err = writeSVR4Header(w.w, hdr)
	return
}

// Write writes to the current file in the CPIO archive. Write returns the error
// ErrWriteTooLong if more than Header.Size bytes are written after WriteHeader.
//
// Calling Write on special types like TypeLink, TypeChar, TypeBlock, TypeDir,
// and TypeFifo returns (0, ErrWriteTooLong) regardless of what the Header.Size
// claims.
func (w *Writer) Write(p []byte) (n int, err error) {
	if w.closed {
		err = ErrWriteAfterClose
		return
	}
	overwrite := false
	if int64(len(p)) > w.nb {
		p = p[0:w.nb]
		overwrite = true
	}
	n, err = w.w.Write(p)
	w.nb -= int64(n)
	if err == nil && overwrite {
		err = ErrWriteTooLong
		return
	}
	w.err = err
	return
}

// Close closes the CPIO archive by flushing the padding, and writing the
// footer. If the current file (from a prior call to WriteHeader) is not fully
// written, then this returns an error.
func (w *Writer) Close() error {
	if w.err != nil || w.closed {
		return w.err
	}
	w.err = w.WriteHeader(trailer)
	if w.err != nil {
		return w.err
	}
	w.Flush()
	w.closed = true
	return w.err
}

// AddFS adds the files from fs.FS to the archive.
//
// It walks the directory tree starting at the root of the filesystem adding
// each file to the tar archive while maintaining the directory structure.
func (w *Writer) AddFS(fsys fs.FS) error {
	return fs.WalkDir(fsys, ".", func(
		name string, d fs.DirEntry, err error,
	) error {
		if err != nil {
			return err
		}
		if name == "." {
			return nil
		}

		info, err := d.Info()
		if err != nil {
			return err
		}

		header, err := FileInfoHeader(info)
		if err != nil {
			return err
		}
		header.Name = name

		if err = w.WriteHeader(header); err != nil {
			return fmt.Errorf("write header: %w", err)
		}

		switch d.Type() {
		case os.ModeDir:
			// Directories do not have a body and fail on [fs.File.Read].
			return nil
		case fs.ModeSymlink:
			target, err := internal.ReadLink(fsys, name)
			if err != nil {
				return err
			}

			_, err = w.Write([]byte(target))

			return err
		case 0:
			file, err := fsys.Open(name)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(w, file)

			return err
		default:
			return fs.ErrInvalid
		}
	})
}