fix(tvix/tests/tvix-init): run as pid1, support init= r/6651

We were previously not running as PID1, but were executing things via
uinitcmd (by u-root init).
This had some problems - a real system pid1 really wants to be pid1.

u-root init itself seems to be too complicated to configure as a
chainloader, so instead, we just provide init/pid1 ourselves and do the
little bit of mount setup / stage 1 that's necessary.

`tvix.exec` is renamed to `tvix.run`, to further distinguish it from the
`exec` bash function.

Our init will now also power off the machine if the shell, or the custom
tvix.run command finished. In the case of switch_root (`init=` being
passed), we terminate before, so it's up to the real system anyways.

Change-Id: If00c0c11ca9393968e76ae0d131594b5c135fa48
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9454
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Tested-by: BuildkiteCI

1 files changed, 75 insertions, 22 deletions

diff --git a/tvix/tests/tvix-init.go b/tvix/tests/tvix-init.go
index a81001159ae2..97a24bab3547 100644
--- a/tvix/tests/tvix-init.go
+++ b/tvix/tests/tvix-init.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"log"
 	"os"
 	"os/exec"
@@ -37,11 +38,41 @@ func parseCmdline(cmdline string) map[string]string {
 	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() {
-	log.Println("Running tvix-init…")
+	fmt.Print(`
+  ______      _         ____      _ __
+ /_  __/   __(_)  __   /  _/___  (_) /_
+  / / | | / / / |/_/   / // __ \/ / __/
+ / /  | |/ / />  <   _/ // / / / / /_
+/_/   |___/_/_/|_|  /___/_/ /_/_/\__/
+
+`)
 
-	log.Println("Creating /nix/store")
-	os.MkdirAll("/nix/store", os.ModePerm)
+	// 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 {
@@ -49,37 +80,59 @@ func main() {
 	}
 	cmdlineFields := parseCmdline(string(cmdline))
 
-	log.Println("Mounting…")
-	if err := run("mount", "-t", "virtiofs", "tvix", "/nix/store", "-o", "ro"); err != nil {
-		log.Printf("Failed to run mount: %v\n", err)
-	}
-
-	// If tvix.find is set, invoke find /nix/store
 	if _, ok := cmdlineFields["tvix.find"]; ok {
-		log.Println("Listing…")
+		// 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 tvix.shell is set, invoke the elvish shell
-	if v, ok := cmdlineFields["tvix.shell"]; ok {
-		log.Printf("Invoking shell%s\n…", v)
 		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)
+		}
 
-	// If tvix.exec is set, invoke the binary specified
-	if v, ok := cmdlineFields["tvix.exec"]; ok {
-		log.Printf("Invoking %s\n…", v)
-		if err := syscall.Exec(v, []string{v}, []string{}); err != nil {
-			log.Printf("Failed to exec: %s\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!")
 	}
 
-	log.Println("Powering off")
+	// 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 command: %v\n", err)
+		log.Printf("Failed to run poweroff command: %v\n", err)
 	}
 }