banger/internal/system/files.go

package system

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"
)

const (
	minWorkSeedBytes   int64 = 512 * 1024 * 1024
	workSeedSlackBytes int64 = 256 * 1024 * 1024
	workSeedRoundBytes int64 = 64 * 1024 * 1024
)

func CopyFilePreferClone(sourcePath, targetPath string) error {
	source, err := os.Open(sourcePath)
	if err != nil {
		return err
	}
	defer source.Close()

	info, err := source.Stat()
	if err != nil {
		return err
	}

	target, err := os.OpenFile(targetPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, info.Mode().Perm())
	if err != nil {
		return err
	}
	defer target.Close()

	if err := unix.IoctlFileClone(int(target.Fd()), int(source.Fd())); err == nil {
		return nil
	}
	if _, err := source.Seek(0, io.SeekStart); err != nil {
		return err
	}
	if _, err := target.Seek(0, io.SeekStart); err != nil {
		return err
	}
	if _, err := io.Copy(target, source); err != nil {
		return err
	}
	if err := target.Sync(); err != nil {
		return err
	}
	if err := target.Chmod(info.Mode().Perm()); err != nil {
		return err
	}
	return nil
}

func WorkSeedPath(rootfsPath string) string {
	rootfsPath = strings.TrimSpace(rootfsPath)
	if rootfsPath == "" {
		return ""
	}
	if strings.HasSuffix(rootfsPath, ".ext4") {
		return strings.TrimSuffix(rootfsPath, ".ext4") + ".work-seed.ext4"
	}
	return rootfsPath + ".work-seed"
}

func BuildWorkSeedImage(ctx context.Context, runner CommandRunner, rootfsPath, outPath string) error {
	rootMount, cleanupRoot, err := MountTempDir(ctx, runner, rootfsPath, true)
	if err != nil {
		return err
	}
	defer cleanupRoot()

	rootHome := filepath.Join(rootMount, "root")
	sizeBytes, err := estimateWorkSeedSize(ctx, runner, rootHome)
	if err != nil {
		return err
	}
	if err := os.RemoveAll(outPath); err != nil && !os.IsNotExist(err) {
		return err
	}
	file, err := os.OpenFile(outPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	if err := file.Close(); err != nil {
		return err
	}
	if err := os.Truncate(outPath, sizeBytes); err != nil {
		return err
	}
	if _, err := runner.Run(ctx, "mkfs.ext4", "-F", outPath); err != nil {
		return err
	}

	workMount, cleanupWork, err := MountTempDir(ctx, runner, outPath, false)
	if err != nil {
		return err
	}
	defer cleanupWork()

	return CopyDirContents(ctx, runner, rootHome, workMount, true)
}

func estimateWorkSeedSize(ctx context.Context, runner CommandRunner, rootHome string) (int64, error) {
	var usedBytes int64
	err := filepath.Walk(rootHome, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Mode().IsRegular() {
			usedBytes += info.Size()
		}
		return nil
	})
	if err != nil {
		if os.IsPermission(err) {
			out, sudoErr := runner.RunSudo(ctx, "du", "-sb", rootHome)
			if sudoErr != nil {
				return 0, fmt.Errorf("%w; sudo du fallback failed: %v", err, sudoErr)
			}
			return roundWorkSeedSize(parseDuSize(out)), nil
		}
		return 0, err
	}
	return roundWorkSeedSize(usedBytes), nil
}

func roundWorkSeedSize(usedBytes int64) int64 {
	sizeBytes := usedBytes*2 + workSeedSlackBytes
	if sizeBytes < minWorkSeedBytes {
		sizeBytes = minWorkSeedBytes
	}
	if rem := sizeBytes % workSeedRoundBytes; rem != 0 {
		sizeBytes += workSeedRoundBytes - rem
	}
	return sizeBytes
}

func parseDuSize(out []byte) int64 {
	fields := strings.Fields(string(out))
	if len(fields) == 0 {
		return 0
	}
	sizeBytes, err := strconv.ParseInt(fields[0], 10, 64)
	if err != nil {
		return 0
	}
	return sizeBytes
}

func ReadNormalizedLines(path string) ([]string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var out []string
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasSuffix(line, "\r") {
			line = strings.TrimSuffix(line, "\r")
		}
		if idx := strings.Index(line, "#"); idx >= 0 {
			line = line[:idx]
		}
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		out = append(out, line)
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("file has no entries: %s", path)
	}
	return out, nil
}