wsh shellintegration (#189)

cmd/server/main-server.go

@@ -19,6 +19,7 @@ import (
 	"github.com/wavetermdev/thenextwave/pkg/blockcontroller"
 	"github.com/wavetermdev/thenextwave/pkg/filestore"
 	"github.com/wavetermdev/thenextwave/pkg/service"
+	"github.com/wavetermdev/thenextwave/pkg/util/shellutil"
 	"github.com/wavetermdev/thenextwave/pkg/wavebase"
 	"github.com/wavetermdev/thenextwave/pkg/wconfig"
 	"github.com/wavetermdev/thenextwave/pkg/web"
@@ -26,6 +27,10 @@ import (
 	"github.com/wavetermdev/thenextwave/pkg/wstore"
 )
 
+// these are set at build time
+var WaveVersion = "0.0.0"
+var BuildTime = "0"
+
 const ReadySignalPidVarName = "WAVETERM_READY_SIGNAL_PID"
 
 var shutdownOnce sync.Once
@@ -81,6 +86,7 @@ func main() {
 	log.SetPrefix("[wavesrv] ")
 	blockcontroller.WshServerFactoryFn = wshserver.MakeWshServer
 	web.WshServerFactoryFn = wshserver.MakeWshServer
+	wavebase.WaveVersion = WaveVersion
 
 	err := service.ValidateServiceMap()
 	if err != nil {
@@ -104,6 +110,7 @@ func main() {
 		}
 	}()
 
+	log.Printf("wave version: %s (%s)\n", WaveVersion, BuildTime)
 	log.Printf("wave home dir: %s\n", wavebase.GetWaveHomeDir())
 	err = filestore.InitFilestore()
 	if err != nil {
@@ -115,13 +122,18 @@ func main() {
 		log.Printf("error initializing wstore: %v\n", err)
 		return
 	}
+	go func() {
+		err := shellutil.InitCustomShellStartupFiles()
+		if err != nil {
+			log.Printf("error initializing wsh and shell-integration files: %v\n", err)
+		}
+	}()
 	err = wstore.EnsureInitialData()
 	if err != nil {
 		log.Printf("error ensuring initial data: %v\n", err)
 		return
 	}
 	installShutdownSignalHandlers()
-
 	go stdinReadWatch()
 	configWatcher()
 	webListener, err := web.MakeTCPListener("web")

pkg/shellexec/shellexec.go

@@ -11,6 +11,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"reflect"
 	"regexp"
 	"strconv"
@@ -192,31 +193,60 @@ func StartRemoteShellProc(termSize TermSize, cmdStr string, cmdOpts CommandOptsT
 	return &ShellProc{Cmd: sessionWrap, Pty: cmdPty, CloseOnce: &sync.Once{}, DoneCh: make(chan any)}, nil
 }
 
+func isZshShell(shellPath string) bool {
+	// get the base path, and then check contains
+	shellBase := filepath.Base(shellPath)
+	return strings.Contains(shellBase, "zsh")
+}
+
+func isBashShell(shellPath string) bool {
+	// get the base path, and then check contains
+	shellBase := filepath.Base(shellPath)
+	return strings.Contains(shellBase, "bash")
+}
+
 func StartShellProc(termSize TermSize, cmdStr string, cmdOpts CommandOptsType) (*ShellProc, error) {
+	shellutil.InitCustomShellStartupFiles()
 	var ecmd *exec.Cmd
 	var shellOpts []string
-	if cmdOpts.Login {
-		shellOpts = append(shellOpts, "-l")
-	}
-	if cmdOpts.Interactive {
-		shellOpts = append(shellOpts, "-i")
-	}
+
+	shellPath := shellutil.DetectLocalShellPath()
 	if cmdStr == "" {
-		shellPath := shellutil.DetectLocalShellPath()
+		if isBashShell(shellPath) {
+			// add --rcfile
+			// cant set -l or -i with --rcfile
+			shellOpts = append(shellOpts, "--rcfile", shellutil.GetBashRcFileOverride())
+		} else {
+			if cmdOpts.Login {
+				shellOpts = append(shellOpts, "-l")
+			}
+			if cmdOpts.Interactive {
+				shellOpts = append(shellOpts, "-i")
+			}
+		}
 		ecmd = exec.Command(shellPath, shellOpts...)
+		ecmd.Env = os.Environ()
+		if isZshShell(shellPath) {
+			shellutil.UpdateCmdEnv(ecmd, map[string]string{"ZDOTDIR": shellutil.GetZshZDotDir()})
+		}
 	} else {
-		shellPath := shellutil.DetectLocalShellPath()
+		if cmdOpts.Login {
+			shellOpts = append(shellOpts, "-l")
+		}
+		if cmdOpts.Interactive {
+			shellOpts = append(shellOpts, "-i")
+		}
 		shellOpts = append(shellOpts, "-c", cmdStr)
 		ecmd = exec.Command(shellPath, shellOpts...)
+		ecmd.Env = os.Environ()
 	}
-	ecmd.Env = os.Environ()
 	if cmdOpts.Cwd != "" {
 		ecmd.Dir = cmdOpts.Cwd
 	}
 	if cwdErr := checkCwd(ecmd.Dir); cwdErr != nil {
 		ecmd.Dir = wavebase.GetHomeDir()
 	}
-	envToAdd := shellutil.WaveshellEnvVars(shellutil.DefaultTermType)
+	envToAdd := shellutil.WaveshellLocalEnvVars(shellutil.DefaultTermType)
 	if os.Getenv("LANG") == "" {
 		envToAdd["LANG"] = wavebase.DetermineLang()
 	}
@@ -250,7 +280,7 @@ func StartShellProc(termSize TermSize, cmdStr string, cmdOpts CommandOptsType) (
 
 func RunSimpleCmdInPty(ecmd *exec.Cmd, termSize TermSize) ([]byte, error) {
 	ecmd.Env = os.Environ()
-	shellutil.UpdateCmdEnv(ecmd, shellutil.WaveshellEnvVars(shellutil.DefaultTermType))
+	shellutil.UpdateCmdEnv(ecmd, shellutil.WaveshellLocalEnvVars(shellutil.DefaultTermType))
 	cmdPty, cmdTty, err := pty.Open()
 	if err != nil {
 		return nil, fmt.Errorf("opening new pty: %w", err)

pkg/util/shellutil/shellutil.go

@@ -5,15 +5,19 @@ package shellutil
 
 import (
 	"context"
+	"fmt"
+	"log"
 	"os"
 	"os/exec"
 	"os/user"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/wavetermdev/thenextwave/pkg/util/utilfn"
 	"github.com/wavetermdev/thenextwave/pkg/wavebase"
 )
 
@@ -27,6 +31,55 @@ var userShellRegexp = regexp.MustCompile(`^UserShell: (.*)$`)
 
 const DefaultShellPath = "/bin/bash"
 
+const WaveAppPathVarName = "WAVETERM_APP_PATH"
+const AppPathBinDir = "bin"
+
+const (
+	ZshIntegrationDir  = "zsh-integration"
+	BashIntegrationDir = "bash-integration"
+	WaveHomeBinDir     = "bin"
+
+	ZshStartup_Zprofile = `
+# Source the original zprofile
+[ -f ~/.zprofile ] && source ~/.zprofile
+`
+
+	ZshStartup_Zshrc = `
+# Source the original zshrc
+[ -f ~/.zshrc ] && source ~/.zshrc
+
+export PATH=$WAVETERM_HOME/bin:$PATH
+`
+
+	ZshStartup_Zlogin = `
+# Source the original zlogin
+[ -f ~/.zlogin ] && source ~/.zlogin
+`
+
+	ZshStartup_Zshenv = `
+[ -f ~/.zshenv ] && source ~/.zshenv
+`
+
+	BashStartup_Bashrc = `
+# Source /etc/profile if it exists
+if [ -f /etc/profile ]; then
+    . /etc/profile
+fi
+
+# Source the first of ~/.bash_profile, ~/.bash_login, or ~/.profile that exists
+if [ -f ~/.bash_profile ]; then
+    . ~/.bash_profile
+elif [ -f ~/.bash_login ]; then
+    . ~/.bash_login
+elif [ -f ~/.profile ]; then
+    . ~/.profile
+fi
+
+set -i
+export PATH=$WAVETERM_HOME/bin:$PATH
+`
+)
+
 func DetectLocalShellPath() string {
 	shellPath := GetMacUserShell()
 	if shellPath == "" {
@@ -48,7 +101,7 @@ func GetMacUserShell() string {
 	return cachedMacUserShell
 }
 
-// dscl . -read /User/[username] UserShell
+// dscl . -read /Users/[username] UserShell
 // defaults to /bin/bash
 func internalMacUserShell() string {
 	osUser, err := user.Current()
@@ -70,13 +123,15 @@ func internalMacUserShell() string {
 	return m[1]
 }
 
-func WaveshellEnvVars(termType string) map[string]string {
+func WaveshellLocalEnvVars(termType string) map[string]string {
 	rtn := make(map[string]string)
 	if termType != "" {
 		rtn["TERM"] = termType
 	}
+	rtn["TERM_PROGRAM"] = "waveterm"
 	rtn["WAVETERM"], _ = os.Executable()
 	rtn["WAVETERM_VERSION"] = wavebase.WaveVersion
+	rtn["WAVETERM_HOME"] = wavebase.GetWaveHomeDir()
 	return rtn
 }
 
@@ -115,3 +170,93 @@ func GetEnvStrKey(envStr string) string {
 	}
 	return envStr[0:eqIdx]
 }
+
+var initStartupFilesOnce = &sync.Once{}
+
+// in a Once block so it can be called multiple times
+// we run it at startup, but also before launching local shells so we know everything is initialized before starting the shell
+func InitCustomShellStartupFiles() error {
+	var err error
+	initStartupFilesOnce.Do(func() {
+		err = initCustomShellStartupFilesInternal()
+	})
+	return err
+}
+
+func GetBashRcFileOverride() string {
+	return filepath.Join(wavebase.GetWaveHomeDir(), BashIntegrationDir, ".bashrc")
+}
+
+func GetZshZDotDir() string {
+	return filepath.Join(wavebase.GetWaveHomeDir(), ZshIntegrationDir)
+}
+
+func initCustomShellStartupFilesInternal() error {
+	log.Printf("initializing wsh and shell startup files\n")
+	waveHome := wavebase.GetWaveHomeDir()
+	zshDir := filepath.Join(waveHome, ZshIntegrationDir)
+	err := wavebase.CacheEnsureDir(zshDir, ZshIntegrationDir, 0755, ZshIntegrationDir)
+	if err != nil {
+		return err
+	}
+	bashDir := filepath.Join(waveHome, BashIntegrationDir)
+	err = wavebase.CacheEnsureDir(bashDir, BashIntegrationDir, 0755, BashIntegrationDir)
+	if err != nil {
+		return err
+	}
+	binDir := filepath.Join(waveHome, WaveHomeBinDir)
+	err = wavebase.CacheEnsureDir(binDir, WaveHomeBinDir, 0755, WaveHomeBinDir)
+	if err != nil {
+		return err
+	}
+
+	zprofilePath := filepath.Join(zshDir, ".zprofile")
+	err = os.WriteFile(zprofilePath, []byte(ZshStartup_Zprofile), 0644)
+	if err != nil {
+		return fmt.Errorf("error writing zsh-integration .zprofile: %v", err)
+	}
+	zshrcPath := filepath.Join(zshDir, ".zshrc")
+	err = os.WriteFile(zshrcPath, []byte(ZshStartup_Zshrc), 0644)
+	if err != nil {
+		return fmt.Errorf("error writing zsh-integration .zshrc: %v", err)
+	}
+	zloginPath := filepath.Join(zshDir, ".zlogin")
+	err = os.WriteFile(zloginPath, []byte(ZshStartup_Zlogin), 0644)
+	if err != nil {
+		return fmt.Errorf("error writing zsh-integration .zlogin: %v", err)
+	}
+	zshenvPath := filepath.Join(zshDir, ".zshenv")
+	err = os.WriteFile(zshenvPath, []byte(ZshStartup_Zshenv), 0644)
+	if err != nil {
+		return fmt.Errorf("error writing zsh-integration .zshenv: %v", err)
+	}
+	bashrcPath := filepath.Join(bashDir, ".bashrc")
+	err = os.WriteFile(bashrcPath, []byte(BashStartup_Bashrc), 0644)
+	if err != nil {
+		return fmt.Errorf("error writing bash-integration .bashrc: %v", err)
+	}
+
+	// copy the correct binary to bin
+	appPath := os.Getenv(WaveAppPathVarName)
+	if appPath == "" {
+		return fmt.Errorf("no app path set")
+	}
+	appBinPath := filepath.Join(appPath, AppPathBinDir)
+	wshBaseName := computeWshBaseName()
+	wshFullPath := filepath.Join(appBinPath, wshBaseName)
+	if _, err := os.Stat(wshFullPath); err != nil {
+		log.Printf("error (non-fatal), could not resolve wsh binary %q: %v\n", wshFullPath, err)
+		return nil
+	}
+	wshDstPath := filepath.Join(binDir, "wsh")
+	err = utilfn.AtomicRenameCopy(wshDstPath, wshFullPath, 0755)
+	if err != nil {
+		return fmt.Errorf("error copying wsh binary to bin: %v", err)
+	}
+	log.Printf("wsh binary successfully %q copied to %q\n", wshBaseName, wshDstPath)
+	return nil
+}
+
+func computeWshBaseName() string {
+	return fmt.Sprintf("wsh-%s-%s.%s", wavebase.WaveVersion, runtime.GOOS, runtime.GOARCH)
+}

pkg/util/utilfn/utilfn.go

@@ -837,3 +837,35 @@ func MergeStrMaps[T any](m1 map[string]T, m2 map[string]T) map[string]T {
 	}
 	return rtn
 }
+
+func AtomicRenameCopy(dstPath string, srcPath string, perms os.FileMode) error {
+	// first copy the file to dstPath.new, then rename into place
+	srcFd, err := os.Open(srcPath)
+	if err != nil {
+		return err
+	}
+	defer srcFd.Close()
+	tempName := dstPath + ".new"
+	dstFd, err := os.Create(tempName)
+	if err != nil {
+		return err
+	}
+	defer dstFd.Close()
+	_, err = io.Copy(dstFd, srcFd)
+	if err != nil {
+		return err
+	}
+	err = dstFd.Close()
+	if err != nil {
+		return err
+	}
+	err = os.Chmod(tempName, perms)
+	if err != nil {
+		return err
+	}
+	err = os.Rename(tempName, dstPath)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/wavebase/wavebase.go

@@ -20,7 +20,9 @@ import (
 	"github.com/alexflint/go-filemutex"
 )
 
-const WaveVersion = "v0.1.0"
+// set by main-server.go
+var WaveVersion = "0.0.0"
+
 const DefaultWaveHome = "~/.w2"
 const DevWaveHome = "~/.w2-dev"
 const WaveHomeVarName = "WAVETERM_HOME"