depot/go/nix/bnix/bnix.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"

	"golang.org/x/crypto/ssh"
	"hg.lukegb.com/lukegb/depot/go/nix/nixdrv"
	"hg.lukegb.com/lukegb/depot/go/nix/nixstore"
)

var (
	remoteFlag = flag.String("remote", "unix://", "Remote.")
)

func ensure(ctx context.Context, d *nixstore.Daemon, at *nixstore.ActivityTracker, path string) error {
	return d.EnsurePath(at, path)
}

var rs nixdrv.Resolver

func loadDerivation(ctx context.Context, path string) (*nixdrv.Derivation, error) {
	return rs.LoadDerivation(path)
}

func buildDerivation(ctx context.Context, d *nixstore.Daemon, at *nixstore.ActivityTracker, path string) error {
	drv, err := loadDerivation(ctx, path)
	if err != nil {
		return fmt.Errorf("loading derivation: %w", err)
	}

	basicDrv, err := drv.ToBasicDerivation(rs)
	if err != nil {
		return fmt.Errorf("resolving: %w", err)
	}

	brs, err := d.BuildDerivation(at, path, basicDrv, nixstore.BMNormal)
	if err != nil {
		return err
	}
	log.Printf("build result: %v", brs)
	return nil
}

func connectTo(ctx context.Context, remote string) (*nixstore.Daemon, nixdrv.Resolver, error) {
	u, err := url.Parse(remote)
	if err != nil {
		return nil, nil, fmt.Errorf("parsing remote %q as URL: %w", remote, err)
	}

	switch u.Scheme {
	case "unix":
		if u.Path == "" {
			u.Path = nixstore.DaemonSock
		}
		d, err := nixstore.OpenDaemon(u.Path)
		if err != nil {
			return nil, nil, err
		}
		return d, nixdrv.LocalFSResolver{}, nil
	case "ssh-ng":
		// Construct a ClientConfig from the URL.
		cfg := &ssh.ClientConfig{}
		if u.Query().Has("privkey") {
			var keys []ssh.Signer
			for _, privkeyPath := range u.Query()["privkey"] {
				privkeyF, err := os.Open(privkeyPath)
				if err != nil {
					return nil, nil, fmt.Errorf("opening privkey %q: %w", privkeyPath, err)
				}
				defer privkeyF.Close()
				privkeyB, err := io.ReadAll(privkeyF)
				if err != nil {
					return nil, nil, fmt.Errorf("reading privkey %q: %w", privkeyPath, err)
				}

				privkey, err := ssh.ParsePrivateKey(privkeyB)
				if err != nil {
					return nil, nil, fmt.Errorf("parsing privkey %q: %w", privkeyPath, err)
				}
				keys = append(keys, privkey)
			}
			cfg.Auth = append(cfg.Auth, ssh.PublicKeys(keys...))
		}
		if u.User != nil {
			cfg.User = u.User.Username()
			if pw, ok := u.User.Password(); ok {
				cfg.Auth = append(cfg.Auth, ssh.Password(pw))
			}
		}
		switch {
		case u.Query().Has("host-key"):
			hkStr := u.Query().Get("host-key")
			_, _, hk, _, _, err := ssh.ParseKnownHosts(append([]byte("x "), []byte(hkStr)...))
			if err != nil {
				return nil, nil, fmt.Errorf("parsing host-key %q: %w", hkStr, err)
			}
			cfg.HostKeyCallback = ssh.FixedHostKey(hk)
		case u.Query().Has("insecure-allow-any-ssh-host-key"):
			cfg.HostKeyCallback = ssh.InsecureIgnoreHostKey()
		default:
			return nil, nil, fmt.Errorf("some SSH host key configuration is required (?host-key=; ?insecure-allow-any-ssh-host-key)")
		}

		// Work out other misc parameters.
		// ...remote command.
		remoteCmd := "nix-daemon --stdio"
		if u.Query().Has("remote-cmd") {
			remoteCmd = u.Query().Get("remote-cmd")
		}

		// Work out the host:port to connect to.
		remote := u.Hostname()
		if portStr := u.Port(); portStr != "" {
			remote = remote + ":" + portStr
		} else {
			remote = remote + ":22"
		}

		conn, err := ssh.Dial("tcp", remote, cfg)
		if err != nil {
			return nil, nil, fmt.Errorf("dialing %v via SSH: %w", remote, err)
		}
		sess, err := conn.NewSession()
		if err != nil {
			conn.Close()
			return nil, nil, fmt.Errorf("opening SSH session to %v: %w", remote, err)
		}
		stdin, err := sess.StdinPipe()
		if err != nil {
			conn.Close()
			return nil, nil, fmt.Errorf("opening stdin pipe: %w", err)
		}
		stdout, err := sess.StdoutPipe()
		if err != nil {
			conn.Close()
			return nil, nil, fmt.Errorf("opening stdout pipe: %w", err)
		}
		if err := sess.Start(remoteCmd); err != nil {
			conn.Close()
			return nil, nil, fmt.Errorf("starting %q: %w", remoteCmd, err)
		}
		d, err := nixstore.OpenDaemonWithIOs(stdout, stdin, conn)
		if err != nil {
			conn.Close()
			return nil, nil, fmt.Errorf("establishing connection to daemon: %w", err)
		}
		return d, nixdrv.LocalFSResolver{}, nil // TODO: change from LocalFSResolver?
	default:
		return nil, nil, fmt.Errorf("unknown remote %q", remote)
	}
}

func main() {
	flag.Parse()

	badCall := func(f string, xs ...interface{}) {
		fmt.Fprintf(os.Stderr, f+"\n", xs...)
		flag.Usage()
		os.Exit(1)
	}

	if flag.NArg() < 1 {
		badCall("need a subcommand")
	}

	var cmd func(context.Context, *nixstore.Daemon, *nixstore.ActivityTracker) error
	switch flag.Arg(0) {
	case "ensure":
		if flag.NArg() != 2 {
			badCall("`ensure` needs a store path")
		}
		cmd = func(ctx context.Context, d *nixstore.Daemon, at *nixstore.ActivityTracker) error {
			return ensure(ctx, d, at, flag.Arg(1))
		}
	case "show-derivation":
		if flag.NArg() != 2 {
			badCall("`show-derivation` needs a derivation")
		}
		cmd = func(ctx context.Context, d *nixstore.Daemon, at *nixstore.ActivityTracker) error {
			drv, err := loadDerivation(ctx, flag.Arg(1))
			if err != nil {
				return err
			}
			fmt.Printf("%#v\n", drv)
			return nil
		}
	case "build-derivation":
		if flag.NArg() != 2 {
			badCall("`build-derivation` needs a derivation")
		}
		cmd = func(ctx context.Context, d *nixstore.Daemon, at *nixstore.ActivityTracker) error {
			return buildDerivation(ctx, d, at, flag.Arg(1))
		}
	default:
		badCall("bad subcommand %s", flag.Arg(0))
	}

	at := nixstore.NewActivityTracker()

	ctx := context.Background()
	d, rss, err := connectTo(ctx, *remoteFlag)
	if err != nil {
		log.Fatalf("connectTo(%q): %v", *remoteFlag, err)
	}
	defer d.Close()
	rs = rss

	if err := cmd(ctx, d, at); err != nil {
		log.Fatalf("%s: %s", flag.Arg(0), err)
	}
}