1 files changed, 138 insertions, 0 deletions
diff --git a/tvix/boot/tvix-init.go b/tvix/boot/tvix-init.go
new file mode 100644
index 000000000000..97a24bab3547
--- /dev/null
+++ b/tvix/boot/tvix-init.go
@@ -0,0 +1,138 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"os/exec"
+	"strings"
+	"syscall"
+)
+
+// run the given command, connecting std{in,err,out} with the OS one.
+func run(args ...string) error {
+	cmd := exec.Command(args[0], args[1:]...)
+	cmd.Stdin = os.Stdin
+	cmd.Stderr = os.Stderr
+	cmd.Stdout = os.Stdout
+
+	return cmd.Run()
+}
+
+// parse the cmdline, return a map[string]string.
+func parseCmdline(cmdline string) map[string]string {
+	line := strings.TrimSuffix(cmdline, "\n")
+	fields := strings.Fields(line)
+	out := make(map[string]string, 0)
+
+	for _, arg := range fields {
+		kv := strings.SplitN(arg, "=", 2)
+		switch len(kv) {
+		case 1:
+			out[kv[0]] = ""
+		case 2:
+			out[kv[0]] = kv[1]
+		}
+	}
+
+	return out
+}
+
+// mounts the nix store from the virtiofs tag to the given destination,
+// creating the destination if it doesn't exist already.
+func mountTvixStore(dest string) error {
+	if err := os.MkdirAll(dest, os.ModePerm); err != nil {
+		return fmt.Errorf("unable to mkdir dest: %w", err)
+	}
+	if err := run("mount", "-t", "virtiofs", "tvix", dest, "-o", "ro"); err != nil {
+		return fmt.Errorf("unable to run mount: %w", err)
+	}
+
+	return nil
+}
+
+func main() {
+	fmt.Print(`
+  ______      _         ____      _ __
+ /_  __/   __(_)  __   /  _/___  (_) /_
+  / / | | / / / |/_/   / // __ \/ / __/
+ / /  | |/ / />  <   _/ // / / / / /_
+/_/   |___/_/_/|_|  /___/_/ /_/_/\__/
+
+`)
+
+	// Set PATH to "/bbin", so we can find the u-root tools
+	os.Setenv("PATH", "/bbin")
+
+	if err := run("mount", "-t", "proc", "none", "/proc"); err != nil {
+		log.Printf("Failed to mount /proc: %v\n", err)
+	}
+	if err := run("mount", "-t", "sysfs", "none", "/sys"); err != nil {
+		log.Printf("Failed to mount /sys: %v\n", err)
+	}
+	if err := run("mount", "-t", "devtmpfs", "devtmpfs", "/dev"); err != nil {
+		log.Printf("Failed to mount /dev: %v\n", err)
+	}
+
+	cmdline, err := os.ReadFile("/proc/cmdline")
+	if err != nil {
+		log.Printf("Failed to read cmdline: %s\n", err)
+	}
+	cmdlineFields := parseCmdline(string(cmdline))
+
+	if _, ok := cmdlineFields["tvix.find"]; ok {
+		// If tvix.find is set, invoke find /nix/store
+		if err := mountTvixStore("/nix/store"); err != nil {
+			log.Printf("Failed to mount tvix store: %v\n", err)
+		}
+
+		if err := run("find", "/nix/store"); err != nil {
+			log.Printf("Failed to run find command: %s\n", err)
+		}
+	} else if _, ok := cmdlineFields["tvix.shell"]; ok {
+		// If tvix.shell is set, mount the nix store to /nix/store directly,
+		// then invoke the elvish shell
+		if err := mountTvixStore("/nix/store"); err != nil {
+			log.Printf("Failed to mount tvix store: %v\n", err)
+		}
+
+		if err := run("elvish"); err != nil {
+			log.Printf("Failed to run shell: %s\n", err)
+		}
+	} else if v, ok := cmdlineFields["tvix.run"]; ok {
+		// If tvix.run is set, mount the nix store to /nix/store directly,
+		// then invoke the command.
+		if err := mountTvixStore("/nix/store"); err != nil {
+			log.Printf("Failed to mount tvix store: %v\n", err)
+		}
+
+		if err := run(v); err != nil {
+			log.Printf("Failed to run command: %s\n", err)
+		}
+	} else if v, ok := cmdlineFields["init"]; ok {
+		// If init is set, invoke the binary specified (with switch_root),
+		// and prepare /fs beforehand as well.
+		os.Mkdir("/fs", os.ModePerm)
+		if err := run("mount", "-t", "tmpfs", "none", "/fs"); err != nil {
+			log.Fatalf("Failed to mount /fs tmpfs: %s\n", err)
+		}
+
+		// Mount /fs/nix/store
+		if err := mountTvixStore("/fs/nix/store"); err != nil {
+			log.Fatalf("Failed to mount tvix store: %v\n", err)
+		}
+
+		// Invoke switch_root, which will take care of moving /proc, /sys and /dev.
+		if err := syscall.Exec("/bbin/switch_root", []string{"switch_root", "/fs", v}, []string{}); err != nil {
+			log.Printf("Failed to switch root: %s\n", err)
+		}
+	} else {
+		log.Printf("No command detected, not knowing what to do!")
+	}
+
+	// This is only reached in the non switch_root case.
+	log.Printf("Nothing left to be done, powering off.")
+	if err := run("poweroff"); err != nil {
+		log.Printf("Failed to run poweroff command: %v\n", err)
+	}
+}