feat(tools/depot-scanner): init project r/1683

depot-scanner is a tool that runs Nix and parses the --trace-file-access output to deduce what files are necessary to evaluate a derivation.

Take DEPOT_ROOT from the environment. If depotRoot doesn't exist, print an error early.

Fix the build of the protobuf library. Switch to the GRPC build rule, as a service is in this proto file.

Create the PathType enum and parse it from cmdline flags.

Change-Id: I537b5c6bceecf76ca510f7ac04ab9dad7785feb1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1769
Tested-by: BuildkiteCI
Reviewed-by: tazjin <mail@tazj.in>
Reviewed-by: lukegb <lukegb@tvl.fyi>

5 files changed, 280 insertions, 0 deletions

diff --git a/tools/depot-scanner/OWNERS b/tools/depot-scanner/OWNERS
new file mode 100644
index 0000000000..cefacea4d0
--- /dev/null
+++ b/tools/depot-scanner/OWNERS
@@ -0,0 +1,3 @@
+inherit: true
+owners:
+ - riking
diff --git a/tools/depot-scanner/default.nix b/tools/depot-scanner/default.nix
new file mode 100644
index 0000000000..d18034cff0
--- /dev/null
+++ b/tools/depot-scanner/default.nix
@@ -0,0 +1,16 @@
+{ depot, pkgs, ...}:
+
+let
+  localProto = depot.nix.buildGo.grpc {
+    name = "code.tvl.fyi/tools/depot-scanner/proto";
+    proto = ./depot_scanner.proto;
+  };
+in depot.nix.buildGo.program {
+  name = "depot-scanner";
+  srcs = [
+    ./main.go
+  ];
+  deps = [
+    localProto
+  ];
+} // { inherit localProto; }
diff --git a/tools/depot-scanner/depot_scanner.proto b/tools/depot-scanner/depot_scanner.proto
new file mode 100644
index 0000000000..5249daebf4
--- /dev/null
+++ b/tools/depot-scanner/depot_scanner.proto
@@ -0,0 +1,46 @@
+// Copyright 2020 TVL
+// SPDX-License-Identifier: MIT
+
+syntax = "proto3";
+package tvl.tools.depot_scanner;
+option go_package = "code.tvl.fyi/tools/depot-scanner/proto";
+
+enum PathType {
+  UNKNOWN = 0;
+  DEPOT = 1;
+  STORE = 2;
+  CORE = 3;
+}
+
+message ScanRequest {
+  // Which revision of the depot
+  string revision = 1;
+  string attr = 2;
+  // Optionally, the attr to evaluate can be provided as a path to a folder or a
+  // .nix file.  This is used by the HTTP service.
+  string attrAsPath = 3;
+}
+
+message ScanResponse {
+  repeated string depotPath = 1;
+  repeated string nixStorePath = 2;
+  repeated string corePkgsPath = 4;
+  repeated string otherPath = 3;
+
+  bytes derivation = 5;
+}
+
+message ArchiveRequest {
+  repeated string depotPath = 1;
+}
+
+message ArchiveChunk {
+  bytes chunk = 1;
+}
+
+service DepotScanService {
+  rpc Scan(ScanRequest) returns (ScanResponse);
+
+  rpc MakeArchive(ArchiveRequest) returns (stream ArchiveChunk);
+}
+
diff --git a/tools/depot-scanner/go.mod b/tools/depot-scanner/go.mod
new file mode 100644
index 0000000000..bdd22fc1ef
--- /dev/null
+++ b/tools/depot-scanner/go.mod
@@ -0,0 +1,3 @@
+module code.tvl.fyi/tools/depot-scanner
+
+go 1.14
diff --git a/tools/depot-scanner/main.go b/tools/depot-scanner/main.go
new file mode 100644
index 0000000000..8d69981c0d
--- /dev/null
+++ b/tools/depot-scanner/main.go
@@ -0,0 +1,212 @@
+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)
+	}
+
+}