From 75c274c10448c89b30fccf699a7956389ea3c664 Mon Sep 17 00:00:00 2001 From: Mike Sawka Date: Wed, 31 Jul 2024 23:47:33 -0700 Subject: [PATCH] wsh shellintegration (#189) --- cmd/server/main-server.go | 14 ++- pkg/shellexec/shellexec.go | 52 ++++++++--- pkg/util/shellutil/shellutil.go | 149 +++++++++++++++++++++++++++++++- pkg/util/utilfn/utilfn.go | 32 +++++++ pkg/wavebase/wavebase.go | 4 +- 5 files changed, 236 insertions(+), 15 deletions(-) diff --git a/cmd/server/main-server.go b/cmd/server/main-server.go index b5579a3a..0759eeef 100644 --- a/cmd/server/main-server.go +++ b/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") diff --git a/pkg/shellexec/shellexec.go b/pkg/shellexec/shellexec.go index 4a9e729a..67a277d0 100644 --- a/pkg/shellexec/shellexec.go +++ b/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) diff --git a/pkg/util/shellutil/shellutil.go b/pkg/util/shellutil/shellutil.go index ce635826..58f53c32 100644 --- a/pkg/util/shellutil/shellutil.go +++ b/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) +} diff --git a/pkg/util/utilfn/utilfn.go b/pkg/util/utilfn/utilfn.go index 7f1aaa9e..04e4644f 100644 --- a/pkg/util/utilfn/utilfn.go +++ b/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 +} diff --git a/pkg/wavebase/wavebase.go b/pkg/wavebase/wavebase.go index 58e8e386..e380ddad 100644 --- a/pkg/wavebase/wavebase.go +++ b/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"