tvl-depot/tools/depot-scanner/main.go
Kane York a9650041bb feat(tools/tvlc): pass list of derivations to depot-scanner
You can now provide a list of Nix derivations to tvlc to get a git worktree + sparse-checkout containing only the paths needed to build the specified derivations.

Known bugs: even though //third_party is only passed to readdir(), git doesn't know this and includes all of //third_party/*.

Change-Id: I9dccebd3fbff4bb04ebd568175cf0a7e37d71ab3
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1826
Tested-by: BuildkiteCI
Reviewed-by: tazjin <mail@tazj.in>
2020-08-25 22:34:02 +00:00

222 lines
5.2 KiB
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")
var relativeFlag = flag.Bool("relpath", false, "when printing paths, print them relative to the root of their path type")
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] {
if *relativePath {
k = strings.TrimPrefix(k, *nixStoreRoot)
k = strings.TrimPrefix(k, "/")
}
fmt.Println(k)
}
}
if enabledPathTypes[pb.PathType_DEPOT] {
for k, _ := range results[depotPath] {
if *relativeFlag {
k = strings.TrimPrefix(k, *depotRoot)
k = strings.TrimPrefix(k, "/")
}
fmt.Println(k)
}
}
if enabledPathTypes[pb.PathType_CORE] {
for k, _ := range results[corePkgsPath] {
// TODO relativeFlag
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)
}
}