nixbuild/cmd/nix-tool/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"slices"
	"syscall"

	"gensokyo.uk/nix"
	"hakurei.app/command"
)

type commandHandlerError string

func (c commandHandlerError) Error() string { return string(c) }

func main() {
	var ctx nix.Context

	nixCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	var (
		flagNixOS   bool
		flagVerbose bool
		flagJSON    bool
	)
	c := command.New(os.Stderr, log.Printf, "nix-tool", func(args []string) error {
		log.SetFlags(0)

		var stderr io.Writer
		if flagVerbose {
			stderr = os.Stderr
		}
		ctx = nix.New(nixCtx, nil, os.Stdout, stderr)

		return nil
	}).
		Flag(&flagNixOS, "nixos", command.BoolFlag(false), "Interpret input as NixOS flake installable").
		Flag(&flagVerbose, "v", command.BoolFlag(false), "Connect nix stderr").
		Flag(&flagJSON, "json", command.BoolFlag(false), "Serialise output in JSON when applicable")

	// this is most likely what you are here for
	var (
		flagCacheKeyPath string

		flagCacheComp     string
		flagCachePComp    bool
		flagCacheBucket   string
		flagCacheEndpoint string
		flagCacheRegion   string
		flagCacheScheme   string
		flagCacheCredPath string
	)
	c.NewCommand("cache", "Build an installable and upload it to a binary cache", func(args []string) error {
		if err := replaceFromEnviron(map[string]*string{
			"NIX_TOOL_CACHE_KEY":         &flagCacheKeyPath,
			"NIX_TOOL_CACHE_COMPRESSION": &flagCacheComp,
			"NIX_TOOL_CACHE_BUCKET":      &flagCacheBucket,
			"NIX_TOOL_CACHE_ENDPOINT":    &flagCacheEndpoint,
			"NIX_TOOL_CACHE_REGION":      &flagCacheRegion,
			"NIX_TOOL_CACHE_SCHEME":      &flagCacheScheme,
			"NIX_TOOL_CACHE_CRED":        &flagCacheCredPath,
		}); err != nil {
			return err
		}

		var (
			installable string
			collective  []string
		)
		if err := buildAndResolve(ctx, "cache", &installable, &collective, flagNixOS, args); err != nil {
			return err
		}

		log.Println("signing collected paths...")
		if err := nix.Sign(ctx, flagCacheKeyPath, slices.Values(collective)); err != nil {
			return commandHandlerError(fmt.Sprintf("cannot sign: %v", err))
		}

		log.Println("copying to binary cache...")
		if err := nix.Copy(ctx, flagCacheKeyPath, &nix.BinaryCache{
			Compression:         flagCacheComp,
			ParallelCompression: flagCachePComp,
			Bucket:              flagCacheBucket,
			Endpoint:            flagCacheEndpoint,
			Region:              flagCacheRegion,
			Scheme:              flagCacheScheme,
			CredentialsPath:     flagCacheCredPath,
		}, slices.Values(collective)); err != nil {
			return commandHandlerError(fmt.Sprintf("cannot copy: %v", err))
		}

		return nil
	}).
		Flag(&flagCacheKeyPath, "secret-key", command.StringFlag("env"), "Path to the secret key used to sign the binary cache").
		Flag(&flagCacheComp, "compression", command.StringFlag("zstd"), "Name of the compression algorithm to use").
		Flag(&flagCachePComp, "parallel-compression", command.BoolFlag(true), "Whether parallel compression is enabled").
		Flag(&flagCacheBucket, "bucket", command.StringFlag("env"), "S3 bucket name").
		Flag(&flagCacheEndpoint, "endpoint", command.StringFlag("env"), "S3 endpoint").
		Flag(&flagCacheRegion, "region", command.StringFlag("env"), "S3 region").
		Flag(&flagCacheScheme, "scheme", command.StringFlag("https"), "S3 scheme").
		Flag(&flagCacheCredPath, "credentials-path", command.StringFlag("env"), "Path to the s3 shared credentials file")

	var (
		resolveFlagOut string
	)
	c.Command("resolve", "Resolve possible build-time dependencies", func(args []string) error {
		var (
			installable string
			collective  []string
		)
		if err := buildAndResolve(ctx, "resolve", &installable, &collective, flagNixOS, args); err != nil {
			return err
		}

		f, err := os.Create(resolveFlagOut)
		if err != nil {
			return commandHandlerError(fmt.Sprintf("cannot create store paths file: %v", err))
		}
		if flagJSON {
			if err = json.NewEncoder(f).Encode(collective); err != nil {
				return commandHandlerError(fmt.Sprintf("cannot serialise JSON: %v", err))
			}
		} else {
			if _, err = nix.WriteStdin(f, slices.Values(collective)); err != nil {
				return commandHandlerError(fmt.Sprintf("cannot write store path list: %v", err))
			}
		}
		return nil
	}).
		Flag(&resolveFlagOut, "o", command.StringFlag("store-paths"), "Path to write collected store paths")

	c.Command("stat", "Compute the instantiated deduplication stat of an installable", func(args []string) error {
		var installable string
		if err := formatInstallable("instantiated", &installable, flagNixOS, args); err != nil {
			return err
		}

		log.Println("initialising evaluator...")
		var collective []string
		if eval, err := nix.NewInstantiatedEvaluator(ctx, installable); err != nil {
			return commandHandlerError(fmt.Sprintf("cannot initialise evaluator: %v", err))
		} else {
			log.Println("collecting paths...")
			collective = slices.Collect(eval.Instantiated())
			if err := eval.Err(); err != nil {
				return commandHandlerError(fmt.Sprintf("cannot collect: %v", err))
			}
		}

		l := len(collective)
		slices.Sort(collective)
		collective = slices.Compact(collective)
		log.Printf("`%s`: %d paths instantiated, %d duplicate, dedup %.1f%%",
			installable, l, l-len(collective), float64(l-len(collective))/float64(l)*100)
		return nil
	})

	c.Command("instantiated", "Evaluate an installable and output all derivations instantiated during evaluation", func(args []string) error {
		var installable string
		if err := formatInstallable("instantiated", &installable, flagNixOS, args); err != nil {
			return err
		}

		if v, err := nix.EvalInstantiated(ctx, installable); err != nil {
			return commandHandlerError(fmt.Sprintf("cannot evaluate for instantiated derivations: %v", err))
		} else if flagJSON {
			if err = json.NewEncoder(os.Stdout).Encode(v); err != nil {
				return commandHandlerError(fmt.Sprintf("cannot serialise JSON: %v", err))
			}
			return nil
		} else {
			_, _ = nix.WriteStdin(os.Stdout, slices.Values(v))
			return nil
		}
	})

	c.MustParse(os.Args[1:], func(err error) {
		var ce commandHandlerError
		if errors.As(err, &ce) {
			log.Fatal(ce.Error())
			return
		}
		log.Fatalf("cannot parse command: %v", err)
	})
}