nix: allow users to set shell startup hooks

config.go

@@ -16,6 +16,12 @@ type Config struct {
 	// Packages is the slice of Nix packages that devbox makes available in
 	// its environment.
 	Packages []string `cue:"[...string]" json:"packages"`
+
+	// Shell configures the devbox shell environment.
+	Shell struct {
+		// InitHook contains commands that will run at shell startup.
+		InitHook string `json:"init_hook,omitempty"`
+	} `json:"shell,omitempty"`
 }
 
 // ReadConfig reads a devbox config file.

devbox.go

@@ -131,7 +131,13 @@ func (d *Devbox) Shell() error {
 		return errors.WithStack(err)
 	}
 	nixDir := filepath.Join(d.srcDir, ".devbox/gen/shell.nix")
-	return nix.Shell(nixDir)
+	sh, err := nix.DetectShell()
+	if err != nil {
+		// Fall back to using a plain Nix shell.
+		sh = &nix.Shell{}
+	}
+	sh.UserInitHook = d.cfg.Shell.InitHook
+	return sh.Run(nixDir)
 }
 
 // saveCfg writes the config file to the devbox directory.

nix/nix.go

@@ -7,94 +7,9 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"os"
 	"os/exec"
-	"strings"
-
-	"github.com/pkg/errors"
-	"go.jetpack.io/devbox/debug"
-	"go.jetpack.io/devbox/shell"
 )
 
-func Shell(path string) error {
-	// nix-shell only runs bash, which isn't great if the user has a
-	// different default shell. Here we try to detect what their current
-	// shell is, and then `exec` it to replace the bash process inside
-	// nix-shell.
-	sh, err := shell.Detect()
-	if err != nil {
-		// Fall back to running the vanilla Nix bash shell.
-		return runFallbackShell(path)
-	}
-
-	// Naively running the user's shell has two problems:
-	//
-	// 1. The shell will source the user's rc file and potentially reorder
-	// the PATH. This is especially a problem with some shims that prepend
-	// their own directories to the front of the PATH, replacing the
-	// Nix-installed packages.
-	// 2. If their shell is bash, we end up double-sourcing their ~/.bashrc.
-	// Once when nix-shell launches bash, and again when we exec it.
-	//
-	// To workaround this, first we store the current (outside of devbox)
-	// PATH in ORIGINAL_PATH. Then we run a "pure" nix-shell to prevent it
-	// from sourcing their ~/.bashrc. From inside the nix-shell (but before
-	// launching the user's preferred shell) we store the PATH again in
-	// PURE_NIX_PATH. When we're finally in the user's preferred shell, we
-	// can use these env vars to set the PATH so that Nix packages are up
-	// front, and all of the other programs come after.
-	//
-	// ORIGINAL_PATH is set by sh.StartCommand.
-	// PURE_NIX_PATH is set by the shell hook in shell.nix.tmpl.
-	sh.PreInitHook = `
-# Update the $PATH so that the user's init script has access to all of their
-# non-devbox programs.
-export PATH="$PURE_NIX_PATH:$ORIGINAL_PATH"
-`
-	sh.PostInitHook = `
-# Update the $PATH again so that the Nix packages take priority over the
-# programs outside of devbox.
-export PATH="$PURE_NIX_PATH:$ORIGINAL_PATH"
-
-# Prepend to the prompt to make it clear we're in a devbox shell.
-export PS1="(devbox) $PS1"
-`
-
-	if debug.IsEnabled() {
-		sh.PostInitHook += `echo "POST-INIT PATH=$PATH"
-`
-	}
-
-	cmd := exec.Command("nix-shell", path)
-	cmd.Args = append(cmd.Args, "--pure", "--command", sh.ExecCommand())
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-
-	debug.Log("Executing nix-shell command: %v", cmd.Args)
-	return errors.WithStack(cmd.Run())
-}
-
-func runFallbackShell(path string) error {
-	cmd := exec.Command("nix-shell", path)
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-
-	debug.Log("Unrecognized user shell, falling back to: %v", cmd.Args)
-	return errors.WithStack(cmd.Run())
-}
-
-func Exec(path string, command []string) error {
-	runCmd := strings.Join(command, " ")
-	cmd := exec.Command("nix-shell", "--run", runCmd)
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	cmd.Dir = path
-	return errors.WithStack(cmd.Run())
-}
-
 func PkgExists(pkg string) bool {
 	_, found := PkgInfo(pkg)
 	return found

nix/shell.go

@@ -0,0 +1,184 @@
+// Copyright 2022 Jetpack Technologies Inc and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package nix
+
+import (
+	"bytes"
+	_ "embed"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"text/template"
+
+	"github.com/pkg/errors"
+	"go.jetpack.io/devbox/debug"
+)
+
+//go:embed shellrc.tmpl
+var shellrcText string
+var shellrcTmpl = template.Must(template.New("shellrc").Parse(shellrcText))
+
+type name string
+
+const (
+	shUnknown name = ""
+	shBash    name = "bash"
+	shZsh     name = "zsh"
+	shKsh     name = "ksh"
+	shPosix   name = "posix"
+)
+
+// Shell configures a user's shell to run in Devbox. Its zero value is a
+// fallback shell that launches a regular Nix shell.
+type Shell struct {
+	name            name
+	binPath         string
+	userShellrcPath string
+
+	// UserInitHook contains commands that will run at shell startup.
+	UserInitHook string
+}
+
+// DetectShell attempts to determine the user's default shell.
+func DetectShell() (*Shell, error) {
+	path := os.Getenv("SHELL")
+	if path == "" {
+		return nil, errors.New("unable to detect the current shell")
+	}
+
+	sh := &Shell{binPath: filepath.Clean(path)}
+	base := filepath.Base(path)
+	// Login shell
+	if base[0] == '-' {
+		base = base[1:]
+	}
+	switch base {
+	case "bash":
+		sh.name = shBash
+		sh.userShellrcPath = rcfilePath(".bashrc")
+	case "zsh":
+		sh.name = shZsh
+		sh.userShellrcPath = rcfilePath(".zshrc")
+	case "ksh":
+		sh.name = shKsh
+		sh.userShellrcPath = rcfilePath(".kshrc")
+	case "dash", "ash", "sh":
+		sh.name = shPosix
+		sh.userShellrcPath = os.Getenv("ENV")
+
+		// Just make up a name if there isn't already an init file set
+		// so we have somewhere to put a new one.
+		if sh.userShellrcPath == "" {
+			sh.userShellrcPath = ".shinit"
+		}
+	default:
+		sh.name = shUnknown
+	}
+	debug.Log("Detected shell: %s", sh.binPath)
+	debug.Log("Recognized shell as: %s", sh.binPath)
+	debug.Log("Looking for user's shell init file at: %s", sh.userShellrcPath)
+	return sh, nil
+}
+
+// rcfilePath returns the absolute path for an rcfile, which is usually in the
+// user's home directory. It doesn't guarantee that the file exists.
+func rcfilePath(basename string) string {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return ""
+	}
+	return filepath.Join(home, basename)
+}
+
+func (s *Shell) Run(nixPath string) error {
+	// Launch a fallback shell if we couldn't find the path to the user's
+	// default shell.
+	if s.binPath == "" {
+		cmd := exec.Command("nix-shell", nixPath)
+		cmd.Stdin = os.Stdin
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+
+		debug.Log("Unrecognized user shell, falling back to: %v", cmd.Args)
+		return errors.WithStack(cmd.Run())
+	}
+
+	cmd := exec.Command("nix-shell", nixPath)
+	cmd.Args = append(cmd.Args, "--pure", "--command", s.execCommand())
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	debug.Log("Executing nix-shell command: %v", cmd.Args)
+	return errors.WithStack(cmd.Run())
+}
+
+// execCommand is a command that replaces the current shell with s.
+func (s *Shell) execCommand() string {
+	shellrc, err := writeDevboxShellrc(s.userShellrcPath, s.UserInitHook)
+	if err != nil {
+		debug.Log("Failed to write devbox shellrc: %v", err)
+		return "exec " + s.binPath
+	}
+
+	switch s.name {
+	case shBash:
+		return fmt.Sprintf(`exec /usr/bin/env ORIGINAL_PATH="%s" %s --rcfile "%s"`,
+			os.Getenv("PATH"), s.binPath, shellrc)
+	case shZsh:
+		return fmt.Sprintf(`exec /usr/bin/env ORIGINAL_PATH="%s" ZDOTDIR="%s" %s`,
+			os.Getenv("PATH"), filepath.Dir(shellrc), s.binPath)
+	case shKsh, shPosix:
+		return fmt.Sprintf(`exec /usr/bin/env ORIGINAL_PATH="%s" ENV="%s" %s `,
+			os.Getenv("PATH"), shellrc, s.binPath)
+	default:
+		return "exec " + s.binPath
+	}
+}
+
+func writeDevboxShellrc(userShellrcPath string, userHook string) (path string, err error) {
+	// We need a temp dir (as opposed to a temp file) because zsh uses
+	// ZDOTDIR to point to a new directory containing the .zshrc.
+	tmp, err := os.MkdirTemp("", "devbox")
+	if err != nil {
+		return "", fmt.Errorf("create temp dir for shell init file: %v", err)
+	}
+
+	// This is a best-effort to include the user's existing shellrc. If we
+	// can't read it, then just omit it from the devbox shellrc.
+	userShellrc, err := os.ReadFile(userShellrcPath)
+	if err != nil {
+		userShellrc = []byte{}
+	}
+
+	path = filepath.Join(tmp, filepath.Base(userShellrcPath))
+	shellrcf, err := os.Create(path)
+	if err != nil {
+		return "", fmt.Errorf("write to shell init file: %v", err)
+	}
+	defer func() {
+		cerr := shellrcf.Close()
+		if err == nil {
+			err = cerr
+		}
+	}()
+
+	err = shellrcTmpl.Execute(shellrcf, struct {
+		OriginalInit     string
+		OriginalInitPath string
+		UserHook         string
+	}{
+		OriginalInit:     string(bytes.TrimSpace(userShellrc)),
+		OriginalInitPath: filepath.Clean(userShellrcPath),
+		UserHook:         strings.TrimSpace(userHook),
+	})
+	if err != nil {
+		return "", fmt.Errorf("execute shellrc template: %v", err)
+	}
+
+	debug.Log("Wrote devbox shellrc to: %s", path)
+	return path, nil
+}

nix/shellrc.tmpl

@@ -0,0 +1,62 @@
+{{- /*
+
+// This template defines the shellrc file that the devbox shell will run at
+// startup.
+//
+// It includes the user's original shellrc, which varies depending on their
+// shell. It will either be ~/.bashrc, ~/.zshrc, a path set in ENV, or something
+// else. It also appends any user-defined shell hooks from devbox.json.
+//
+// Devbox also needs to ensure that the shell's PATH, prompt, and a few other
+// things are set correctly at startup. To do this, it must run some commands
+// before and after the user's shellrc. These commands are in the
+// "Devbox Pre/Post-init Hook" sections.
+//
+// The devbox pre/post-init hooks assume two environment variables are already
+// set:
+//
+// - ORIGINAL_PATH - embedded into the command built by Shell.execCommand. It
+//                   preserves the PATH at the time `devbox shell` is invoked.
+// - PURE_NIX_PATH - set by the shell hook in shell.nix.tmpl. It preserves the
+//                   PATH set by Nix's "pure" shell mode.
+
+ */ -}}
+
+# Begin Devbox Pre-init Hook
+
+# Update the $PATH so that the user's init script has access to all of their
+# non-devbox programs.
+export PATH="$PURE_NIX_PATH:$ORIGINAL_PATH"
+
+# End Devbox Pre-init Hook
+
+{{- if .OriginalInit }}
+
+# Begin {{ .OriginalInitPath }}
+
+{{ .OriginalInit }}
+
+# End {{ .OriginalInitPath }}
+
+{{- end }}
+
+# Begin Devbox Post-init Hook
+
+# Update the $PATH again so that the Nix packages take priority over the
+# programs outside of devbox.
+export PATH="$PURE_NIX_PATH:$ORIGINAL_PATH"
+
+# Prepend to the prompt to make it clear we're in a devbox shell.
+export PS1="(devbox) $PS1"
+
+# End Devbox Post-init Hook
+
+{{- if .UserHook }}
+
+# Begin Devbox User Hook
+
+{{ .UserHook }}
+
+# End Devbox User Hook
+
+{{- end }}