sfdupes/scan.go

package main

import (
	"bufio"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
)

// chunk is the number of bytes hashed from each end of a file.
const chunk = 1024

// workQueueDepth bounds the job and result channels feeding the stat
// and hash worker pools.
const workQueueDepth = 1024

// errNotRegular reports a path that stopped being a regular file
// between the walk and stat passes.
var errNotRegular = errors.New("no longer a regular file")

// fileRec carries one file between the stat and hash passes.
type fileRec struct {
	path  string
	size  int64
	mtime int64
}

// runScan implements the scan subcommand: three sequential passes (walk,
// stat, hash) over the tree under -root, emitting one NUL-terminated
// record per regular file on stdout.
func runScan(args []string) {
	fl := flag.NewFlagSet("scan", flag.ExitOnError)
	root := fl.String("root", "/srv", "directory tree to scan")
	workers := fl.Int("workers", runtime.NumCPU(),
		"concurrent workers for the stat and hash passes")
	// The flag set uses ExitOnError, so Parse cannot return a non-nil
	// error; the check keeps the error handled explicitly.
	err := fl.Parse(args)
	if err != nil {
		os.Exit(exitUsage)
	}

	if fl.NArg() > 0 {
		fmt.Fprintf(os.Stderr, "scan: unexpected argument %q\n", fl.Arg(0))
		os.Exit(exitUsage)
	}

	if *workers < 1 {
		*workers = 1
	}

	st, err := os.Stat(*root)
	if err != nil {
		fatalf("root %s: %v", *root, err)
	}

	if !st.IsDir() {
		fatalf("root %s: not a directory", *root)
	}

	paths, walkErrs := walkPass(*root)
	recs, statErrs := statPass(paths, *workers)
	emitted, hashErrs := hashPass(recs, *workers)

	skipped := walkErrs + statErrs + hashErrs
	fmt.Fprintf(os.Stderr, "scan: %d files emitted, %d skipped\n",
		emitted, skipped)
}

// walkPass enumerates every regular file under root. It never follows
// symlinks, never descends into directories named .zfs, and warns and
// continues on any per-path error.
func walkPass(root string) ([]string, int) {
	var (
		paths []string
		errs  int
	)

	prog := newProgress("walk", -1)
	walkErr := filepath.WalkDir(root, func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			errs++

			prog.warnf("walk %s: %v", p, err)

			if d != nil && d.IsDir() {
				return filepath.SkipDir
			}

			return nil
		}

		if d.IsDir() {
			// ZFS snapshot pseudo-dirs would list every file
			// once per snapshot; never descend.
			if d.Name() == ".zfs" {
				return filepath.SkipDir
			}

			return nil
		}
		// Regular files only: skip symlinks, sockets, FIFOs, and
		// device nodes.
		if !d.Type().IsRegular() {
			return nil
		}

		paths = append(paths, p)

		prog.increment()

		return nil
	})

	prog.finish()

	if walkErr != nil {
		fatalf("walk %s: %v", root, walkErr)
	}

	return paths, errs
}

// statPass lstats every collected path in a worker pool, recording size
// and mtime. Paths that fail to stat (or are no longer regular files)
// are warned about and dropped.
func statPass(paths []string, workers int) ([]fileRec, int) {
	type result struct {
		rec fileRec
		err error
	}

	jobs := make(chan string, workQueueDepth)
	results := make(chan result, workQueueDepth)

	for range workers {
		go func() {
			for p := range jobs {
				fi, err := os.Lstat(p)
				switch {
				case err != nil:
					results <- result{rec: fileRec{path: p}, err: err}
				case !fi.Mode().IsRegular():
					results <- result{
						rec: fileRec{path: p},
						err: errNotRegular,
					}
				default:
					results <- result{rec: fileRec{
						path:  p,
						size:  fi.Size(),
						mtime: fi.ModTime().Unix(),
					}}
				}
			}
		}()
	}

	go func() {
		for _, p := range paths {
			jobs <- p
		}

		close(jobs)
	}()

	prog := newProgress("stat", int64(len(paths)))

	var errs int

	recs := make([]fileRec, 0, len(paths))
	for range paths {
		r := <-results
		if r.err != nil {
			errs++

			prog.warnf("stat %s: %v", r.rec.path, r.err)
		} else {
			recs = append(recs, r.rec)
		}

		prog.increment()
	}

	prog.finish()

	return recs, errs
}

// hashResult carries one file's head/tail hashes (or the error that
// prevented hashing it) from the hash workers to the main goroutine.
type hashResult struct {
	rec  fileRec
	head string
	tail string
	err  error
}

// startHashWorkers starts the hash worker pool over recs and returns
// the channel its results arrive on (one per record, in completion
// order).
func startHashWorkers(recs []fileRec, workers int) <-chan hashResult {
	jobs := make(chan fileRec, workQueueDepth)
	results := make(chan hashResult, workQueueDepth)

	for range workers {
		go func() {
			for rec := range jobs {
				head, tail, err := hashHeadTail(rec.path, rec.size)
				results <- hashResult{rec: rec, head: head, tail: tail, err: err}
			}
		}()
	}

	go func() {
		for _, rec := range recs {
			jobs <- rec
		}

		close(jobs)
	}()

	return results
}

// hashPass hashes the first and last chunk bytes of every file in a
// worker pool and emits the output records on stdout from the main
// goroutine. Files that fail to open or read are warned about and
// dropped.
func hashPass(recs []fileRec, workers int) (int, int) {
	results := startHashWorkers(recs, workers)
	prog := newProgress("hash", int64(len(recs)))
	out := bufio.NewWriterSize(os.Stdout, ioBufSize)

	var emitted, errs int

	for range recs {
		r := <-results
		if r.err != nil {
			errs++

			prog.warnf("hash %s: %v", r.rec.path, r.err)
			prog.increment()

			continue
		}

		_, err := fmt.Fprintf(out, "%d\t%d\t%s\t%s\t%s\x00",
			r.rec.size, r.rec.mtime, r.head, r.tail, r.rec.path)
		if err != nil {
			fatalf("write stdout: %v", err)
		}

		emitted++

		prog.increment()
	}

	prog.finish()

	err := out.Flush()
	if err != nil {
		fatalf("write stdout: %v", err)
	}

	return emitted, errs
}

// hashHeadTail returns the lowercase-hex SHA-256 of the first
// min(chunk, size) bytes and of the last min(chunk, size) bytes of the
// file at path. The two reads overlap when size < 2*chunk; for
// size == 0 both hashes are of the empty input. size is the value
// recorded by the stat pass.
func hashHeadTail(path string, size int64) (string, string, error) {
	//nolint:gosec // hashing operator-supplied paths is the tool's purpose
	f, err := os.Open(path)
	if err != nil {
		return "", "", err
	}

	defer func() { _ = f.Close() }()

	n := min(int64(chunk), size)

	buf := make([]byte, n)
	if n > 0 {
		_, err = f.ReadAt(buf, 0)
		if err != nil {
			return "", "", err
		}
	}

	h := sha256.Sum256(buf)

	if n > 0 {
		_, err = f.ReadAt(buf, size-n)
		if err != nil {
			return "", "", err
		}
	}

	t := sha256.Sum256(buf)

	return hex.EncodeToString(h[:]), hex.EncodeToString(t[:]), nil
}