tvl-depot/tools/depot-scanner/main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"

	pb "code.tvl.fyi/tools/depot-scanner/proto"
)

var nixInstantiatePath = flag.String("nix-bin", "/run/current-system/sw/bin/nix-instantiate", "path to nix-instantiate")
var depotRoot = flag.String("depot", envOr("DEPOT_ROOT", "/depot/"), "path to tvl.fyi depot at current canon")
var nixStoreRoot = flag.String("store-path", "/nix/store/", "prefix for all valid nix store paths")

var modeFlag = flag.String("mode", modeArchive, "operation mode. valid values: tar, print")
var onlyFlag = flag.String("only", "", "only enable the listed output types, comma separated. valid values: DEPOT, STORE, CORE, UNKNOWN")

const (
	modeArchive = "tar"
	modePrint   = "print"
)

const (
	// String that identifies a path as belonging to nix corepkgs.
	corePkgsString = "/share/nix/corepkgs/"

	depotTraceString = "trace: depot-scan: "
)

type fileScanType int

const (
	unknownPath fileScanType = iota
	depotPath
	nixStorePath
	corePkgsPath
)

func launchNix(attr string) (*exec.Cmd, io.ReadCloser, io.ReadCloser, error) {
	cmd := exec.Command(*nixInstantiatePath, "--trace-file-access", "-A", attr)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		stdout.Close()
		return nil, nil, nil, err
	}

	err = cmd.Start()
	if err != nil {
		stdout.Close()
		stderr.Close()
		return nil, nil, nil, err
	}

	return cmd, stdout, stderr, nil
}

func categorizePath(path string) fileScanType {
	if strings.HasPrefix(path, *nixStoreRoot) {
		if strings.Contains(path, corePkgsString) {
			return corePkgsPath
		}
		return nixStorePath
	} else if strings.HasPrefix(path, *depotRoot) {
		return depotPath
	} else if strings.Contains(path, corePkgsString) {
		return corePkgsPath
	}
	return unknownPath
}

func addPath(path string, out map[fileScanType]map[string]struct{}) {
	cat := categorizePath(path)
	if out[cat] == nil {
		out[cat] = make(map[string]struct{})
	}

	out[cat][path] = struct{}{}
}

func consumeOutput(stdout, stderr io.ReadCloser) (map[fileScanType]map[string]struct{}, string, error) {
	result := make(map[fileScanType]map[string]struct{})

	scanner := bufio.NewScanner(stderr)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, depotTraceString) {
			addPath(strings.TrimPrefix(line, depotTraceString), result)
		}
	}
	if scanner.Err() != nil {
		return nil, "", scanner.Err()
	}

	// Get derivation path
	derivPath := ""
	scanner = bufio.NewScanner(stdout)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, *nixStoreRoot) {
			derivPath = line
			// consume the rest of the output
		}
	}
	if scanner.Err() != nil {
		return nil, "", scanner.Err()
	}

	return result, derivPath, nil
}

func main() {
	flag.Parse()

	checkDepotRoot()

	enabledPathTypes := make(map[pb.PathType]bool, 4)
	if len(*onlyFlag) > 0 {
		enabledOutputs := strings.Split(*onlyFlag, ",")
		for _, v := range enabledOutputs {
			i, ok := pb.PathType_value[strings.ToUpper(v)]
			if !ok {
				fmt.Fprintln(os.Stderr, "warning: unrecognized PathType name: ", v)
				continue
			}
			enabledPathTypes[pb.PathType(i)] = true
		}
	} else {
		// Default
		enabledPathTypes = map[pb.PathType]bool{
			pb.PathType_UNKNOWN: true,
			pb.PathType_DEPOT:   true,
			pb.PathType_STORE:   true,
			pb.PathType_CORE:    true,
		}
	}

	cmd, stdout, stderr, err := launchNix(flag.Arg(0))
	if err != nil {
		panic(fmt.Errorf("could not launch nix: %w", err))
	}
	results, derivPath, err := consumeOutput(stdout, stderr)
	if err != nil {
		err2 := cmd.Wait()
		if err2 != nil {
			panic(fmt.Errorf("nix-instantiate failed: %w\nadditionally, while reading output: %w", err2, err))
		}
		panic(fmt.Errorf("problem reading nix output: %w", err))
	}
	err = cmd.Wait()
	if err != nil {
		panic(fmt.Errorf("nix-instantiate failed: %w", err))
	}

	_ = derivPath

	if *modeFlag == "print" {
		if enabledPathTypes[pb.PathType_STORE] {
			for k, _ := range results[nixStorePath] {
				fmt.Println(k)
			}
		}
		if enabledPathTypes[pb.PathType_DEPOT] {
			for k, _ := range results[depotPath] {
				fmt.Println(k)
			}
		}
		if enabledPathTypes[pb.PathType_CORE] {
			for k, _ := range results[corePkgsPath] {
				fmt.Println(k)
			}
		}
		if enabledPathTypes[pb.PathType_UNKNOWN] {
			for k, _ := range results[unknownPath] {
				fmt.Println(k)
			}
		}
	} else {
		panic("unimplemented")
	}
}

func envOr(envVar, def string) string {
	v := os.Getenv(envVar)
	if v == "" {
		return def
	}
	return v
}

func checkDepotRoot() {
	if *depotRoot == "" {
		fmt.Fprintln(os.Stderr, "error: DEPOT_ROOT / -depot not set")
		os.Exit(2)
	}
	_, err := os.Stat(*depotRoot)
	if os.IsNotExist(err) {
		fmt.Fprintf(os.Stderr, "error: %q does not exist\ndid you forget to set DEPOT_ROOT / --depot ?\n", *depotRoot)
		os.Exit(1)
	} else if err != nil {
		fmt.Fprintf(os.Stderr, "error: could not stat %q: %v\n", *depotRoot, err)
		os.Exit(1)
	}

}