sfdupes/scan.go

package main

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

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

// 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")
	fl.Parse(args)
	if fl.NArg() > 0 {
		fmt.Fprintf(os.Stderr, "scan: unexpected argument %q\n", fl.Arg(0))
		os.Exit(2)
	}
	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) (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) (recs []fileRec, errs int) {
	type result struct {
		rec fileRec
		err error
	}
	jobs := make(chan string, 1024)
	results := make(chan result, 1024)
	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: fmt.Errorf("no longer a regular file"),
					}
				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)))
	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
}

// 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) (emitted, errs int) {
	type result struct {
		rec  fileRec
		head string
		tail string
		err  error
	}
	jobs := make(chan fileRec, 1024)
	results := make(chan result, 1024)
	for range workers {
		go func() {
			for rec := range jobs {
				head, tail, err := hashHeadTail(rec.path, rec.size)
				results <- result{rec: rec, head: head, tail: tail, err: err}
			}
		}()
	}
	go func() {
		for _, rec := range recs {
			jobs <- rec
		}
		close(jobs)
	}()

	prog := newProgress("hash", int64(len(recs)))
	out := bufio.NewWriterSize(os.Stdout, 1<<20)
	for range recs {
		r := <-results
		if r.err != nil {
			errs++
			prog.warnf("hash %s: %v", r.rec.path, r.err)
			prog.increment()
			continue
		}
		if _, 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); err != nil {
			fatalf("write stdout: %v", err)
		}
		emitted++
		prog.increment()
	}
	prog.finish()
	if err := out.Flush(); 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) (head, tail string, err error) {
	f, err := os.Open(path)
	if err != nil {
		return "", "", err
	}
	defer f.Close()

	n := min(int64(chunk), size)
	buf := make([]byte, n)
	if n > 0 {
		if _, err := f.ReadAt(buf, 0); err != nil {
			return "", "", err
		}
	}
	h := sha256.Sum256(buf)
	if n > 0 {
		if _, err := f.ReadAt(buf, size-n); err != nil {
			return "", "", err
		}
	}
	t := sha256.Sum256(buf)
	return hex.EncodeToString(h[:]), hex.EncodeToString(t[:]), nil
}