dotfiles/patches/0001-Add-Nix-integration.patch

From ce5fb5686e7d54f51dc15aeb1ef4ec08d5635740 Mon Sep 17 00:00:00 2001
From: main <magic_rb@redalder.org>
Date: Tue, 25 Oct 2022 17:22:50 +0200
Subject: [PATCH] Add Nix integration

Signed-off-by: main <magic_rb@redalder.org>
---
 drivers/docker/config.go |  19 ++++++-
 drivers/docker/driver.go |  52 +++++++++++++++++--
 drivers/docker/nix.go    | 109 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 176 insertions(+), 4 deletions(-)
 create mode 100644 drivers/docker/nix.go

diff --git a/drivers/docker/config.go b/drivers/docker/config.go
index 40d98966a..9aa26458f 100644
--- a/drivers/docker/config.go
+++ b/drivers/docker/config.go
@@ -108,6 +108,10 @@ func PluginLoader(opts map[string]string) (map[string]interface{}, error) {
 		conf["nvidia_runtime"] = v
 	}
 
+	if v, ok := opts["docker.gcroots_dir"]; ok {
+		conf["gcroots_dir"] = v
+	}
+
 	return conf, nil
 }
 
@@ -281,6 +285,11 @@ var (
 			hclspec.NewLiteral(`"5m"`),
 		),
 
+		"gcroots_dir": hclspec.NewDefault(
+			hclspec.NewAttr("gcroots_dir", "string", false),
+			hclspec.NewLiteral(`"/nix/var/nix/gcroots/nomad-docker"`),
+		),
+
 		// the duration that the driver will wait for activity from the Docker engine during an image pull
 		// before canceling the request
 		"pull_activity_timeout": hclspec.NewDefault(
@@ -327,7 +336,7 @@ var (
 	// taskConfigSpec is the hcl specification for the driver config section of
 	// a task within a job. It is returned in the TaskConfigSchema RPC
 	taskConfigSpec = hclspec.NewObject(map[string]*hclspec.Spec{
-		"image":                  hclspec.NewAttr("image", "string", true),
+		"image":                  hclspec.NewAttr("image", "string", false),
 		"advertise_ipv6_address": hclspec.NewAttr("advertise_ipv6_address", "bool", false),
 		"args":                   hclspec.NewAttr("args", "list(string)", false),
 		"auth": hclspec.NewBlock("auth", false, hclspec.NewObject(map[string]*hclspec.Spec{
@@ -402,6 +411,10 @@ var (
 		"volumes":         hclspec.NewAttr("volumes", "list(string)", false),
 		"volume_driver":   hclspec.NewAttr("volume_driver", "string", false),
 		"work_dir":        hclspec.NewAttr("work_dir", "string", false),
+
+		"nix_flake_ref":        hclspec.NewAttr("nix_flake_ref", "string", false),
+		"nix_flake_sha":        hclspec.NewAttr("nix_flake_sha", "string", false),
+		"nix_flake_store_path": hclspec.NewAttr("nix_flake_store_path", "string", false),
 	})
 
 	// driverCapabilities represents the RPC response for what features are
@@ -474,6 +486,10 @@ type TaskConfig struct {
 	VolumeDriver      string             `codec:"volume_driver"`
 	WorkDir           string             `codec:"work_dir"`
 
+	NixFlakeRef       string             `codec:"nix_flake_ref"`
+	NixFlakeSha       string             `codec:"nix_flake_sha"`
+	NixFlakeStorePath string             `codec:"nix_flake_store_path"`
+
 	// MountsList supports the pre-1.0 mounts array syntax
 	MountsList []DockerMount `codec:"mounts"`
 }
@@ -642,6 +657,8 @@ type DriverConfig struct {
 	ExtraLabels                   []string      `codec:"extra_labels"`
 	Logging                       LoggingConfig `codec:"logging"`
 
+	GCRootsDir                    string        `codec:"gcroots_dir"`
+
 	AllowRuntimesList []string            `codec:"allow_runtimes"`
 	allowRuntimes     map[string]struct{} `codec:"-"`
 }
diff --git a/drivers/docker/driver.go b/drivers/docker/driver.go
index 0aa993845..812952fad 100644
--- a/drivers/docker/driver.go
+++ b/drivers/docker/driver.go
@@ -14,6 +14,7 @@ import (
 	"strings"
 	"sync"
 	"time"
+	"os/exec"
 
 	docker "github.com/fsouza/go-dockerclient"
 	"github.com/hashicorp/consul-template/signals"
@@ -254,7 +255,7 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 		return nil, nil, fmt.Errorf("failed to decode driver config: %v", err)
 	}
 
-	if driverConfig.Image == "" {
+	if driverConfig.Image == "" && !(driverConfig.NixFlakeRef != "" && driverConfig.NixFlakeSha != "" && driverConfig.NixFlakeStorePath != "") {
 		return nil, nil, fmt.Errorf("image name required for docker driver")
 	}
 
@@ -269,6 +270,79 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 		return nil, nil, fmt.Errorf("Failed to connect to docker daemon: %s", err)
 	}
 
+	if driverConfig.NixFlakeRef != "" && driverConfig.NixFlakeSha != "" {
+		driverConfig.Image = "magicrb/nix-container-base@sha256:01f199486f5b0e3c90411d700436395f21154f8234b6dfa86eb224eb5b6ad43b";
+
+		nixExecutable, err := exec.LookPath("nix")
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to find `nix` executable")
+		}
+
+		if _, err := os.Stat(driverConfig.NixFlakeStorePath); err != nil {
+			err = NixBuildFlake(nixExecutable, driverConfig.NixFlakeRef, driverConfig.NixFlakeSha)
+			if err != nil {
+				return nil, nil, err
+			}
+
+
+			deps, err := NixGetDeps(nixExecutable, driverConfig.NixFlakeRef)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			for _, dep := range deps {
+				var mount DockerMount
+				mount.Type = "bind"
+				mount.Target = dep;
+				mount.Source = dep;
+				mount.ReadOnly = true;
+
+				driverConfig.Mounts = append(driverConfig.Mounts, mount);
+			}
+
+			storePath, err := NixGetStorePath(nixExecutable, driverConfig.NixFlakeRef)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			driverConfig.Entrypoint[0] = storePath + "/" + driverConfig.Entrypoint[0]
+
+			os.Symlink(storePath, GetGCRoot(d.config.GCRootsDir, cfg.Name, cfg.AllocID))
+		} else {
+			d.eventer.EmitEvent(&drivers.TaskEvent{
+				TaskID:    cfg.ID,
+				AllocID:   cfg.AllocID,
+				TaskName:  cfg.Name,
+				Timestamp: time.Now(),
+				Message:   "Skipping nix build as store path exists",
+				Annotations: map[string]string{
+					"store_path": driverConfig.NixFlakeStorePath,
+				},
+			})
+			deps, err := NixGetDeps(nixExecutable, driverConfig.NixFlakeStorePath)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			for _, dep := range deps {
+				var mount DockerMount
+				mount.Type = "bind"
+				mount.Target = dep;
+				mount.Source = dep;
+				mount.ReadOnly = true;
+
+				driverConfig.Mounts = append(driverConfig.Mounts, mount);
+			}
+
+			driverConfig.Entrypoint[0] = driverConfig.NixFlakeStorePath + "/" + driverConfig.Entrypoint[0]
+
+			os.Symlink(driverConfig.NixFlakeStorePath, GetGCRoot(d.config.GCRootsDir, cfg.Name, cfg.AllocID))
+		}
+	}
+	if (driverConfig.NixFlakeRef != "") != (driverConfig.NixFlakeSha != "") {
+		d.logger.Warn("one of either nix_flake_ref or nix_flake_sha is not set", "container_id", cfg.ID, "nix_flake_ref", driverConfig.NixFlakeRef, "nix_flake_sha", driverConfig.NixFlakeSha)
+	}
+
 	id, err := d.createImage(cfg, &driverConfig, client)
 	if err != nil {
 		return nil, nil, err
@@ -1263,7 +1305,7 @@ func (d *Driver) toDockerMount(m *DockerMount, task *drivers.TaskConfig) (*docke
 
 		// paths inside alloc dir are always allowed as they mount within
 		// a container, and treated as relative to task dir
-		if !d.config.Volumes.Enabled && !isParentPath(task.AllocDir, hm.Source) {
+		if !d.config.Volumes.Enabled && !isParentPath(task.AllocDir, hm.Source) && !isParentPath("/nix/store", hm.Source) {
 			return nil, fmt.Errorf(
 				"volumes are not enabled; cannot mount host path: %q %q",
 				hm.Source, task.AllocDir)
@@ -1425,7 +1467,11 @@ func (d *Driver) StopTask(taskID string, timeout time.Duration, signal string) e
 		return drivers.ErrTaskNotFound
 	}
 
-	return h.Kill(timeout, signal)
+	err := h.Kill(timeout, signal)
+
+	os.Remove(GetGCRoot(d.config.GCRootsDir, h.task.Name, h.task.AllocID))
+
+	return err
 }
 
 func (d *Driver) DestroyTask(taskID string, force bool) error {
diff --git a/drivers/docker/nix.go b/drivers/docker/nix.go
new file mode 100644
index 000000000..426cc51fd
--- /dev/null
+++ b/drivers/docker/nix.go
@@ -0,0 +1,109 @@
+package docker
+
+import (
+	"fmt"
+	"os/exec"
+	"strings"
+	"encoding/json"
+)
+
+func NixGetDeps(executable string, flakeRef string) ([]string, error) {
+	nixDepsCmd := &exec.Cmd {
+		Path: executable,
+		Args: []string{
+			executable,
+			"path-info",
+			"-r",
+			flakeRef,
+		},
+	}
+	res, err := nixDepsCmd.Output()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dependencies of built flake-ref %s", flakeRef)
+	}
+	deps := strings.Split(strings.Trim(string(res), " \n"), "\n")
+
+	return deps, nil
+}
+
+func NixBuildFlake(executable string, flakeRef string, flakeSha string) error {
+
+	flakeHost := strings.Split(flakeRef, "#")
+
+	if len(flakeHost) != 2 {
+		return fmt.Errorf("Invalid flake ref.")
+	}
+
+	nixShaCmd := &exec.Cmd {
+		Path: executable,
+		Args: []string{
+			executable,
+			"flake",
+			"metadata",
+			"--json",
+			flakeHost[0],
+		},
+	}
+	nixSha, err := nixShaCmd.Output()
+	if err != nil {
+		return fmt.Errorf("failed to get sha for flake-ref %s with %s:\n %s", flakeRef, err, string(nixSha))
+	}
+
+	var shaJson map[string]interface{}
+	err = json.Unmarshal(nixSha, &shaJson)
+
+	if err != nil {
+		return fmt.Errorf("failed to parse json %s", err)
+	}
+
+	lockedVal, ok := shaJson["locked"].(map[string]interface{})
+	if !ok {
+		return fmt.Errorf("failed to parse `nix flake metadata` output")
+	}
+	fetchedSha, ok := lockedVal["narHash"].(string)
+	if !ok {
+		return fmt.Errorf("failed to parse `nix flake metadata` output")
+	}
+
+	if string(fetchedSha) != flakeSha {
+		return fmt.Errorf("pinned flake sha doesn't match: \"%s\" != \"%s\"", flakeSha, fetchedSha)
+	}
+
+	nixBuildCmd := &exec.Cmd {
+		Path: executable,
+		Args: []string{
+			executable,
+			"build",
+			"--no-link",
+			flakeRef,
+		},
+	}
+	res, err := nixBuildCmd.Output()
+	if err != nil {
+		return fmt.Errorf("failed to build flake-ref %s with %s:\n %s", flakeRef, err, string(res))
+	}
+
+	return nil
+}
+
+func NixGetStorePath(executable string, flakeRef string) (string, error) {
+	nixEvalCmd := exec.Cmd {
+		Path: executable,
+		Args: []string{
+			executable,
+			"eval",
+			"--raw",
+			flakeRef + ".outPath",
+		},
+	}
+
+	storePath, err := nixEvalCmd.Output()
+	if err != nil {
+		return "", fmt.Errorf("failed to get store path of %s", flakeRef)
+	}
+	return string(storePath), nil
+}
+
+func GetGCRoot(gcRootsDir string, containerName string, allocationId string) string {
+	return fmt.Sprintf("%s/%s-%s", gcRootsDir, containerName, allocationId)
+}
-- 
2.37.1