Ability to run Nix jobs

6 files changed, 301 insertions, 109 deletions

diff --git a/example/agent.hcl b/example/agent.hcl
index c704243..740c221 100644
--- a/example/agent.hcl
+++ b/example/agent.hcl
@@ -3,10 +3,7 @@
 client {
 }
 
-plugin "exec2-driver" {
+plugin "nix2-driver" {
   config {
-    bind_read_only = {
-      "/etc" = "/etc",
-    }
   }
 }
diff --git a/example/example.hcl b/example/example.hcl
index 80da070..dee0e0e 100644
--- a/example/example.hcl
+++ b/example/example.hcl
@@ -3,53 +3,20 @@ job "example" {
   type        = "batch"
 
   group "example" {
-    task "test-host-bin" {
-      driver = "exec2"
-
-      config {
-        command =  "/bin/sh"
-        args = ["-c", "echo hello world"]
-        bind_read_only = {
-          "/bin" = "/bin",
-          "/lib" = "/lib",
-          "/lib64" = "/lib64",
-          "/usr" = "/usr",
-          "/nix" = "/nix",
-        }
-      }
-      user = "lx"
-    }
-
     task "test-nix-hello" {
-      driver = "exec2"
+      driver = "nix2"
 
       config {
-        command = "/sw/bin/nix"
+        command = "sh"
         args = [
-          "--extra-experimental-features", "flakes",
-          "--extra-experimental-features", "nix-command",
-          "run",
+          "-c",
+          "pwd; ls -l *; mount; hello"
+        ]
+        packages = [
+          "github:NixOS/nixpkgs#coreutils",
+          "github:NixOS/nixpkgs#bash",
           "github:NixOS/nixpkgs#hello"
           ]
-        bind = {
-          "/nix" = "/nix",
-        }
-        bind_read_only = {
-          "/home/lx/.nix-profile" = "/sw",
-        }
-      }
-      user = "lx"
-    }
-
-    task "test-nix-store" {
-      driver = "exec2"
-
-      config {
-        command =  "/nix/store/30j23057fqnnc1p4jqmq73p0gxgn0frq-bash-5.1-p16/bin/sh"
-        args = ["-c", "/nix/store/y41s1vcn0irn9ahn9wh62yx2cygs7qjj-coreutils-8.32/bin/ls /*; /nix/store/y41s1vcn0irn9ahn9wh62yx2cygs7qjj-coreutils-8.32/bin/id"]
-        bind_read_only = {
-          "/nix" = "/nix",
-        }
       }
       user = "lx"
     }
diff --git a/example/example2.hcl b/example/example2.hcl
new file mode 100644
index 0000000..8b56f8a
--- /dev/null
+++ b/example/example2.hcl
@@ -0,0 +1,28 @@
+job "example2" {
+  datacenters = ["dc1"]
+  type        = "service"
+
+  group "example" {
+    task "server" {
+      driver = "nix2"
+
+      config {
+        packages = [
+          "github:nixos/nixpkgs#python3",
+          "github:nixos/nixpkgs#bash",
+          "github:nixos/nixpkgs#coreutils",
+          "github:nixos/nixpkgs#curl",
+          "github:nixos/nixpkgs#nix",
+          "github:nixos/nixpkgs#git",
+          "github:nixos/nixpkgs#cacert",
+          "github:nixos/nixpkgs#strace",
+          "github:nixos/nixpkgs#gnugrep",
+          "github:nixos/nixpkgs#mount",
+        ]
+        command = "python3"
+        args = [ "-m", "http.server", "8080" ]
+      }
+      user = "lx"
+    }
+  }
+}
diff --git a/executor/executor_linux.go b/executor/executor_linux.go
index 8665fd0..19bead8 100644
--- a/executor/executor_linux.go
+++ b/executor/executor_linux.go
@@ -801,18 +801,41 @@ func cmdMounts(mounts []*drivers.MountConfig) []*lconfigs.Mount {
 //
 // See also executor.lookupBin for a version used by non-isolated drivers.
 func lookupTaskBin(command *ExecCommand) (string, string, error) {
+	cmd := command.Cmd
+
+	taskPath, hostPath, err := lookupBinFile(command, cmd)
+	if err == nil {
+		return taskPath, hostPath, nil
+	}
+
+	if !strings.Contains(cmd, "/") {
+		// Look up also in /bin
+		bin := filepath.Join("/bin", cmd)
+		taskPath, hostPath, err = lookupBinFile(command, bin)
+		if err == nil {
+			return taskPath, hostPath, nil
+		}
+
+		return "", "", fmt.Errorf("file %s not found in task dir or in mounts, even when looking up /bin", cmd)
+	} else {
+		// If there's a / in the binary's path, we can't fallback to a PATH search
+		return "", "", fmt.Errorf("file %s not found in task dir or in mounts", cmd)
+	}
+
+}
+
+func lookupBinFile(command *ExecCommand, bin string) (string, string, error) {
 	taskDir := command.TaskDir
-	bin := command.Cmd
 
 	// Check in the local directory
 	localDir := filepath.Join(taskDir, allocdir.TaskLocal)
-	taskPath, hostPath, err := getPathInTaskDir(command.TaskDir, localDir, bin)
+	taskPath, hostPath, err := getPathInTaskDir(taskDir, localDir, bin)
 	if err == nil {
 		return taskPath, hostPath, nil
 	}
 
 	// Check at the root of the task's directory
-	taskPath, hostPath, err = getPathInTaskDir(command.TaskDir, command.TaskDir, bin)
+	taskPath, hostPath, err = getPathInTaskDir(taskDir, taskDir, bin)
 	if err == nil {
 		return taskPath, hostPath, nil
 	}
@@ -825,31 +848,7 @@ func lookupTaskBin(command *ExecCommand) (string, string, error) {
 		}
 	}
 
-	// If there's a / in the binary's path, we can't fallback to a PATH search
-	if strings.Contains(bin, "/") {
-		return "", "", fmt.Errorf("file %s not found under path %s", bin, taskDir)
-	}
-
-	// look for a file using a PATH-style lookup inside the directory
-	// root. Similar to the stdlib's exec.LookPath except:
-	//   - uses a restricted lookup PATH rather than the agent process's PATH env var.
-	//   - does not require that the file is already executable (this will be ensured
-	//     by the caller)
-	//   - does not prevent using relative path as added to exec.LookPath in go1.19
-	//     (this gets fixed-up in the caller)
-
-	// This is a fake PATH so that we're not using the agent's PATH
-	restrictedPaths := []string{"/usr/local/bin", "/usr/bin", "/bin"}
-
-	for _, dir := range restrictedPaths {
-		pathDir := filepath.Join(command.TaskDir, dir)
-		taskPath, hostPath, err = getPathInTaskDir(command.TaskDir, pathDir, bin)
-		if err == nil {
-			return taskPath, hostPath, nil
-		}
-	}
-
-	return "", "", fmt.Errorf("file %s not found under path", bin)
+	return "", "", fmt.Errorf("file %s not found in task dir or in mounts", bin)
 }
 
 // getPathInTaskDir searches for the binary in the task directory and nested
diff --git a/nix2/driver.go b/nix2/driver.go
index c97efc5..833e515 100644
--- a/nix2/driver.go
+++ b/nix2/driver.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"sync"
 	"time"
 
@@ -15,7 +16,6 @@ import (
 	"github.com/hashicorp/nomad/client/lib/cgutil"
 	"github.com/hashicorp/nomad/drivers/shared/capabilities"
 	"github.com/hashicorp/nomad/drivers/shared/eventer"
-	"github.com/hashicorp/nomad/drivers/shared/resolvconf"
 	"github.com/hashicorp/nomad/helper/pluginutils/hclutils"
 	"github.com/hashicorp/nomad/helper/pluginutils/loader"
 	"github.com/hashicorp/nomad/helper/pointer"
@@ -89,6 +89,7 @@ var (
 		"ipc_mode":       hclspec.NewAttr("ipc_mode", "string", false),
 		"cap_add":        hclspec.NewAttr("cap_add", "list(string)", false),
 		"cap_drop":       hclspec.NewAttr("cap_drop", "list(string)", false),
+		"packages":       hclspec.NewAttr("packages", "list(string)", false),
 	})
 
 	// driverCapabilities represents the RPC response for what features are
@@ -208,9 +209,12 @@ type TaskConfig struct {
 
 	// CapDrop is a set of linux capabilities to disable.
 	CapDrop []string `codec:"cap_drop"`
+
+	// List of Nix packages to add to environment
+	Packages []string `codec:"packages"`
 }
 
-func (tc *TaskConfig) validate() error {
+func (tc *TaskConfig) validate(dc *Config) error {
 	switch tc.ModePID {
 	case "", executor.IsolationModePrivate, executor.IsolationModeHost:
 	default:
@@ -233,6 +237,12 @@ func (tc *TaskConfig) validate() error {
 		return fmt.Errorf("cap_drop configured with capabilities not supported by system: %s", badDrops)
 	}
 
+	if !dc.AllowBind {
+		if len(tc.Bind) > 0 || len(tc.BindReadOnly) > 0 {
+			return fmt.Errorf("bind and bind_read_only are deactivated for the %s driver", pluginName)
+		}
+	}
+
 	return nil
 }
 
@@ -447,8 +457,14 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 	if err := cfg.DecodeDriverConfig(&driverConfig); err != nil {
 		return nil, nil, fmt.Errorf("failed to decode driver config: %v", err)
 	}
+	if driverConfig.Bind == nil {
+		driverConfig.Bind = make(hclutils.MapStrStr)
+	}
+	if driverConfig.BindReadOnly == nil {
+		driverConfig.BindReadOnly = make(hclutils.MapStrStr)
+	}
 
-	if err := driverConfig.validate(); err != nil {
+	if err := driverConfig.validate(&d.config); err != nil {
 		return nil, nil, fmt.Errorf("failed driver config validation: %v", err)
 	}
 
@@ -475,52 +491,73 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 		user = "0"
 	}
 
-	if cfg.DNS != nil {
-		dnsMount, err := resolvconf.GenerateDNSMount(cfg.TaskDir().Dir, cfg.DNS)
-		if err != nil {
-			return nil, nil, fmt.Errorf("failed to build mount for resolv.conf: %v", err)
+	// Prepare NixOS packages and setup a bunch of read-only mounts
+	// for system stuff and required NixOS packages
+	d.eventer.EmitEvent(&drivers.TaskEvent{
+		TaskID:    cfg.ID,
+		AllocID:   cfg.AllocID,
+		TaskName:  cfg.Name,
+		Timestamp: time.Now(),
+		Message:   "Building Nix packages and preparing NixOS state",
+		Annotations: map[string]string{
+			"packages": strings.Join(driverConfig.Packages, " "),
+		},
+	})
+	taskDirs := cfg.TaskDir()
+	systemMounts, err := prepareNixPackages(taskDirs.Dir, driverConfig.Packages)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Some files are necessary and should be taken from outside if not present already
+	for _, f := range []string{ "/etc/resolv.conf", "/etc/passwd", "/etc/nsswitch.conf" } {
+		if _, ok := systemMounts[f]; !ok {
+			systemMounts[f] = f
 		}
-		cfg.Mounts = append(cfg.Mounts, dnsMount)
 	}
 
-	// Bind mounts specified in driver config
+	d.logger.Info("adding RO system mounts for Nix stuff / system stuff", "system_mounts", hclog.Fmt("%+v", systemMounts))
 
-	// Bind mounts specified in task config
-	if d.config.AllowBind {
-		if driverConfig.Bind != nil {
-			for host, task := range driverConfig.Bind {
-				mount_config := drivers.MountConfig{
-					TaskPath:        task,
-					HostPath:        host,
-					Readonly:        false,
-					PropagationMode: "private",
-				}
-				d.logger.Info("adding RW mount from task spec", "mount_config", hclog.Fmt("%+v", mount_config))
-				cfg.Mounts = append(cfg.Mounts, &mount_config)
-			}
+	for host, task := range systemMounts {
+		mount_config := drivers.MountConfig{
+			TaskPath:        task,
+			HostPath:        host,
+			Readonly:        true,
+			PropagationMode: "private",
 		}
-		if driverConfig.BindReadOnly != nil {
-			for host, task := range driverConfig.BindReadOnly {
-				mount_config := drivers.MountConfig{
-					TaskPath:        task,
-					HostPath:        host,
-					Readonly:        true,
-					PropagationMode: "private",
-				}
-				d.logger.Info("adding RO mount from task spec", "mount_config", hclog.Fmt("%+v", mount_config))
-				cfg.Mounts = append(cfg.Mounts, &mount_config)
-			}
+		cfg.Mounts = append(cfg.Mounts, &mount_config)
+	}
+
+	// Set PATH to /bin
+	cfg.Env["PATH"] = "/bin"
+
+	// Bind mounts specified in task config
+	for host, task := range driverConfig.Bind {
+		mount_config := drivers.MountConfig{
+			TaskPath:        task,
+			HostPath:        host,
+			Readonly:        false,
+			PropagationMode: "private",
 		}
-	} else {
-		if len(driverConfig.Bind) > 0 || len(driverConfig.BindReadOnly) > 0 {
-			return nil, nil, fmt.Errorf("bind and bind_read_only are deactivated for the %s driver", pluginName)
+		d.logger.Info("adding RW mount from task spec", "mount_config", hclog.Fmt("%+v", mount_config))
+		cfg.Mounts = append(cfg.Mounts, &mount_config)
+	}
+	for host, task := range driverConfig.BindReadOnly {
+		mount_config := drivers.MountConfig{
+			TaskPath:        task,
+			HostPath:        host,
+			Readonly:        true,
+			PropagationMode: "private",
 		}
+		d.logger.Info("adding RO mount from task spec", "mount_config", hclog.Fmt("%+v", mount_config))
+		cfg.Mounts = append(cfg.Mounts, &mount_config)
 	}
 
 	caps, err := capabilities.Calculate(
 		capabilities.NomadDefaults(), d.config.AllowCaps, driverConfig.CapAdd, driverConfig.CapDrop,
 	)
 	if err != nil {
+		pluginClient.Kill()
 		return nil, nil, err
 	}
 	d.logger.Debug("task capabilities", "capabilities", caps)
diff --git a/nix2/nix.go b/nix2/nix.go
new file mode 100644
index 0000000..7a86934
--- /dev/null
+++ b/nix2/nix.go
@@ -0,0 +1,164 @@
+package nix2
+
+import (
+	"bytes"
+	"path/filepath"
+	"encoding/json"
+	"fmt"
+	"os"
+	"os/exec"
+
+	"github.com/hashicorp/nomad/helper/pluginutils/hclutils"
+)
+
+const (
+	closureNix = `
+{ path }:
+let
+  nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixos-22.05";
+  inherit (nixpkgs.legacyPackages.x86_64-linux) buildPackages;
+in buildPackages.closureInfo { rootPaths = builtins.storePath path; }
+`
+)
+
+func prepareNixPackages(taskDir string, packages []string) (hclutils.MapStrStr, error) {
+	mounts := make(hclutils.MapStrStr)
+
+	profileLink := filepath.Join(taskDir, "current-profile")
+	profile, err := nixBuildProfile(packages, profileLink)
+	if err != nil {
+		return nil, fmt.Errorf("Build of the flakes failed: %v", err)
+	}
+
+	closureLink := filepath.Join(taskDir, "current-closure")
+	closure, err := nixBuildClosure(profileLink, closureLink)
+	if err != nil {
+		return nil, fmt.Errorf("Build of the flakes failed: %v", err)
+	}
+
+	mounts[profile] = profile
+
+	if entries, err := os.ReadDir(profile); err != nil {
+		return nil, fmt.Errorf("Couldn't read profile directory: %w", err)
+	} else {
+		for _, entry := range entries {
+			if name := entry.Name(); name != "etc" {
+				mounts[filepath.Join(profile, name)] = "/" + name
+				continue
+			}
+
+			etcEntries, err := os.ReadDir(filepath.Join(profile, "etc"))
+			if err != nil {
+				return nil, fmt.Errorf("Couldn't read profile's /etc directory: %w", err)
+			}
+
+			for _, etcEntry := range etcEntries {
+				etcName := etcEntry.Name()
+				mounts[filepath.Join(profile, "etc", etcName)] = "/etc/" + etcName
+			}
+		}
+	}
+
+	mounts[filepath.Join(closure, "registration")] = "/registration"
+
+	requisites, err := nixRequisites(closure)
+	if err != nil {
+		return nil, fmt.Errorf("Couldn't determine flake requisites: %v", err)
+	}
+
+	for _, requisite := range requisites {
+		mounts[requisite] = requisite
+	}
+
+	return mounts, nil
+}
+
+func nixBuildProfile(flakes []string, link string) (string, error) {
+	cmd := exec.Command("nix", append(
+		[]string{
+			"--extra-experimental-features", "nix-command",
+			"--extra-experimental-features", "flakes",
+			"profile",
+			"install",
+			"--no-write-lock-file",
+			"--profile",
+			link},
+		flakes...)...)
+	stderr := &bytes.Buffer{}
+	cmd.Stderr = stderr
+
+	if err := cmd.Run(); err != nil {
+		return "", fmt.Errorf("%v failed: %s. Err: %v", cmd.Args, stderr.String(), err)
+	}
+
+	if target, err := os.Readlink(link); err == nil {
+		return os.Readlink(filepath.Join(filepath.Dir(link), target))
+	} else {
+		return "", err
+	}
+}
+
+func nixBuildClosure(profile string, link string) (string, error) {
+	cmd := exec.Command(
+		"nix",
+		"--extra-experimental-features", "nix-command",
+		"--extra-experimental-features", "flakes",
+		"build",
+		"--out-link", link,
+		"--expr", closureNix,
+		"--impure",
+		"--no-write-lock-file",
+		"--argstr", "path", profile)
+
+	stderr := &bytes.Buffer{}
+	cmd.Stderr = stderr
+
+	if err := cmd.Run(); err != nil {
+		return "", fmt.Errorf("%v failed: %s. Err: %v", cmd.Args, stderr.String(), err)
+	}
+
+	return os.Readlink(link)
+}
+
+type nixPathInfo struct {
+	Path             string   `json:"path"`
+	NarHash          string   `json:"narHash"`
+	NarSize          uint64   `json:"narSize"`
+	References       []string `json:"references"`
+	Deriver          string   `json:"deriver"`
+	RegistrationTime uint64   `json:"registrationTime"`
+	Signatures       []string `json:"signatures"`
+}
+
+func nixRequisites(path string) ([]string, error) {
+	cmd := exec.Command(
+		"nix",
+		"--extra-experimental-features", "nix-command",
+		"--extra-experimental-features", "flakes",
+		"path-info",
+		"--json",
+		"--recursive",
+		path)
+
+	stdout := &bytes.Buffer{}
+	cmd.Stdout = stdout
+
+	stderr := &bytes.Buffer{}
+	cmd.Stderr = stderr
+
+	if err := cmd.Run(); err != nil {
+		return nil, fmt.Errorf("%v failed: %s. Err: %v", cmd.Args, stderr.String(), err)
+	}
+
+	result := []*nixPathInfo{}
+	if err := json.Unmarshal(stdout.Bytes(), &result); err != nil {
+		return nil, err
+	}
+
+	requisites := []string{}
+	for _, result := range result {
+		requisites = append(requisites, result.Path)
+	}
+
+	return requisites, nil
+}