From e5d9547ebdde194b6046be7946b0e451c341bf5b Mon Sep 17 00:00:00 2001
From: Edgar Lee <122112154+elpdt852@users.noreply.github.com>
Date: Wed, 14 Feb 2024 07:09:48 -0500
Subject: [PATCH] Add NewExternalNixBuilder to enable config-level changing of
 NixBuilder

---
 main.go                             |  7 ++++-
 modules/nixos/tests/snapshotter.nix | 43 +++++++++++++++++++++++++----
 pkg/config/config.go                |  7 +++--
 pkg/nix/nix.go                      | 14 ++++++++++
 pkg/plugin/plugin.go                |  8 +++++-
 5 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index b4fdb03..650486f 100644
--- a/main.go
+++ b/main.go
@@ -131,7 +131,12 @@ func serve(ctx context.Context, cfg *config.Config) error {
 		runtime.RegisterImageServiceServer(rpc, imageService)
 	}
 
-	sn, err := nix.NewSnapshotter(cfg.Root)
+	var snapshotterOpts []nix.SnapshotterOpt
+	if cfg.ExternalBuilder != "" {
+		snapshotterOpts = append(snapshotterOpts, nix.WithNixBuilder(nix.NewExternalBuilder(cfg.ExternalBuilder)))
+	}
+
+	sn, err := nix.NewSnapshotter(cfg.Root, snapshotterOpts...)
 	if err != nil {
 		return err
 	}
diff --git a/modules/nixos/tests/snapshotter.nix b/modules/nixos/tests/snapshotter.nix
index 5789157..178da53 100644
--- a/modules/nixos/tests/snapshotter.nix
+++ b/modules/nixos/tests/snapshotter.nix
@@ -1,4 +1,4 @@
-{ lib, ... }:
+{ pkgs, lib, ... }:
 let
   registryHost = "127.0.0.1";
 
@@ -50,8 +50,11 @@ let
       base
       ../nix-snapshotter.nix
     ];
-    services.nix-snapshotter.setContainerdSnapshotter = true;
-    services.nix-snapshotter.enable = true;
+
+    services.nix-snapshotter = {
+      enable = true;
+      setContainerdSnapshotter = true;
+    };
   };
 
   rootless = {
@@ -60,8 +63,11 @@ let
       ../nix-snapshotter.nix
       ../nix-snapshotter-rootless.nix
     ];
-    services.nix-snapshotter.rootless.setContainerdSnapshotter = true;
-    services.nix-snapshotter.rootless.enable = true;
+
+    services.nix-snapshotter.rootless = {
+      enable = true;
+      setContainerdSnapshotter = true;
+    };
 
     users.users.alice = {
       uid = 1000;
@@ -71,8 +77,30 @@ let
 
   both = { imports = [ rootful rootless ]; };
 
+  external = {
+    imports = [
+      base
+      ../nix-snapshotter.nix
+    ];
+
+    services.nix-snapshotter = {
+      enable = true;
+      setContainerdSnapshotter = true;
+      settings.external_builder = pkgs.writeScript "external-builder.sh" ''
+        ${pkgs.nix}/bin/nix build --out-link $1 $2
+      '';
+    };
+  };
+
 in {
-  nodes = { inherit rootful rootless both; };
+  nodes = {
+    inherit 
+      rootful
+      rootless
+      both
+      external
+    ;
+  };
 
   testScript = { nodes, ... }:
     let
@@ -150,5 +178,8 @@ in {
       setup(both)
       test_rootful(both, "both")
       test_rootless(both, "both")
+
+      setup(external)
+      test_rootful(external)
     '';
 }
diff --git a/pkg/config/config.go b/pkg/config/config.go
index db18591..efb0a4a 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -19,9 +19,10 @@ var (
 
 // Config provides nix-snapshotter configuration data.
 type Config struct {
-	Address      string             `toml:"address"`
-	Root         string             `toml:"root"`
-	ImageService ImageServiceConfig `toml:"image_service"`
+	Address         string             `toml:"address"`
+	Root            string             `toml:"root"`
+	ExternalBuilder string             `toml:"external_builder"`
+	ImageService    ImageServiceConfig `toml:"image_service"`
 }
 
 type ImageServiceConfig struct {
diff --git a/pkg/nix/nix.go b/pkg/nix/nix.go
index b29cb06..4554d60 100644
--- a/pkg/nix/nix.go
+++ b/pkg/nix/nix.go
@@ -74,3 +74,17 @@ func defaultNixBuilder(ctx context.Context, outLink, nixStorePath string) error
 	}
 	return err
 }
+
+// NewExternalBuilder returns a NixBuilder from an external executable with
+// two arguments: an out-link path, and a Nix store path.
+func NewExternalBuilder(name string) NixBuilder {
+	return func(ctx context.Context, outLink, nixStorePath string) error {
+		out, err := exec.Command(name, outLink, nixStorePath).CombinedOutput()
+		if err != nil {
+			log.G(ctx).
+				WithField("nixStorePath", nixStorePath).
+				Errorf("Failed to run external nix builder: %s\n%s", err, string(out))
+		}
+		return err
+	}
+}
diff --git a/pkg/plugin/plugin.go b/pkg/plugin/plugin.go
index 1d9454e..c64e996 100644
--- a/pkg/plugin/plugin.go
+++ b/pkg/plugin/plugin.go
@@ -77,7 +77,13 @@ func init() {
 			}
 
 			ic.Meta.Exports["root"] = root
-			return nix.NewSnapshotter(root)
+
+			var snapshotterOpts []nix.SnapshotterOpt
+			if cfg.ExternalBuilder != "" {
+				snapshotterOpts = append(snapshotterOpts, nix.WithNixBuilder(nix.NewExternalBuilder(cfg.ExternalBuilder)))
+			}
+
+			return nix.NewSnapshotter(root, snapshotterOpts...)
 		},
 	})
 }