From 713acb664d98331489a05855482279cb1b0b4074 Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Sat, 20 Jan 2024 13:17:43 -0800
Subject: [PATCH] fix setting of readonly vars issue with zsh rc files

---
 waveshell/pkg/shellapi/zshapi.go      | 34 +++++++++++++++++++++++----
 waveshell/pkg/shellapi/zshapi_test.go | 18 ++++++++++++++
 waveshell/pkg/shexec/shexec.go        |  3 ++-
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/waveshell/pkg/shellapi/zshapi.go b/waveshell/pkg/shellapi/zshapi.go
index 495d35fd..2b4cdf84 100644
--- a/waveshell/pkg/shellapi/zshapi.go
+++ b/waveshell/pkg/shellapi/zshapi.go
@@ -11,6 +11,7 @@ import (
 	"os/exec"
 	"strings"
 	"sync"
+	"unicode"
 
 	"github.com/alessio/shellescape"
 	"github.com/wavetermdev/waveterm/waveshell/pkg/base"
@@ -105,6 +106,7 @@ var localZshMajorVersion = ""
 
 // sentinel value for functions that should be autoloaded
 const ZshFnAutoLoad = "autoload"
+const ZshAutoloadFnBody = "builtin autoload -XU"
 
 type ZshParamKey struct {
 	// paramtype cannot contain spaces
@@ -204,6 +206,19 @@ func makeZshTypesetStmt(varDecl *shellenv.DeclareDeclType) string {
 	}
 }
 
+func isZshSafeNameStr(name string) bool {
+	for _, ch := range name {
+		if ch == '_' {
+			continue
+		}
+		if unicode.IsLetter(ch) || unicode.IsDigit(ch) {
+			continue
+		}
+		return false
+	}
+	return true
+}
+
 func (z zshShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 	var rcBuf bytes.Buffer
 	rcBuf.WriteString(z.GetBaseShellOpts() + "\n")
@@ -216,9 +231,6 @@ func (z zshShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 		if ZshIgnoreVars[varDecl.Name] {
 			continue
 		}
-		if varDecl.IsReadOnly() {
-			continue
-		}
 		if ZshUniqueArrayVars[varDecl.Name] && !varDecl.IsUniqueArray() {
 			varDecl.AddFlag("U")
 		}
@@ -230,7 +242,19 @@ func (z zshShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 		if stmt == "" {
 			continue
 		}
-		rcBuf.WriteString(makeZshTypesetStmt(varDecl))
+		if varDecl.IsReadOnly() {
+			// we can't reset read-only variables
+			// so we check if it is a "safe" name, and then we can write a conditional
+			//     that only sets it if it hasn't already been set.
+			if !isZshSafeNameStr(varDecl.Name) {
+				continue
+			}
+			rcBuf.WriteString(fmt.Sprintf("if (( ! ${+%s} )); then\n", varDecl.Name))
+			rcBuf.WriteString(makeZshTypesetStmt(varDecl))
+			rcBuf.WriteString("\nfi")
+		} else {
+			rcBuf.WriteString(makeZshTypesetStmt(varDecl))
+		}
 		rcBuf.WriteString("\n")
 	}
 	if shellenv.FindVarDecl(varDecls, "ZDOTDIR") == nil {
@@ -519,6 +543,8 @@ func ParseZshFunctions(fpathArr []string, fnBytes []byte, partSeparator []byte)
 		source := fnSource[fnKey.ParamName]
 		if isSourceFileInFpath(fpathArr, source) {
 			fnBody[fnKey] = ZshFnAutoLoad
+		} else if strings.TrimSpace(fnBody[fnKey]) == ZshAutoloadFnBody {
+			fnBody[fnKey] = ZshFnAutoLoad
 		}
 	}
 	return fnBody
diff --git a/waveshell/pkg/shellapi/zshapi_test.go b/waveshell/pkg/shellapi/zshapi_test.go
index e93a0f97..557ac99e 100644
--- a/waveshell/pkg/shellapi/zshapi_test.go
+++ b/waveshell/pkg/shellapi/zshapi_test.go
@@ -27,3 +27,21 @@ func TestParseZshDecl(t *testing.T) {
 	declStr = `typeset -x -g -aT FOO foo=( 1 2 3 )`
 	testSingleDecl(declStr)
 }
+
+func TestZshSafeDeclName(t *testing.T) {
+	if !isZshSafeNameStr("foo") {
+		t.Errorf("foo should be safe")
+	}
+	if isZshSafeNameStr("foo bar") {
+		t.Errorf("foo bar should not be safe")
+	}
+	if !isZshSafeNameStr("foo_bar") {
+		t.Errorf("foo_bar should be safe")
+	}
+	if !isZshSafeNameStr("été") {
+		t.Errorf("été should be be safe")
+	}
+	if isZshSafeNameStr("hello\x01z") {
+		t.Errorf("should not be safe")
+	}
+}
diff --git a/waveshell/pkg/shexec/shexec.go b/waveshell/pkg/shexec/shexec.go
index 956fd181..8ce9f10e 100644
--- a/waveshell/pkg/shexec/shexec.go
+++ b/waveshell/pkg/shexec/shexec.go
@@ -50,6 +50,7 @@ const ShellVarName = "SHELL"
 const SigKillWaitTime = 2 * time.Second
 const RtnStateFdNum = 20
 const ReturnStateReadWaitTime = 2 * time.Second
+const ForceDebugRcFile = false
 
 const ClientCommandFmt = `
 PATH=$PATH:~/.mshell;
@@ -821,7 +822,7 @@ func RunCommandSimple(pk *packet.RunPacketType, sender *packet.PacketSender, fro
 		rcFileStr += trapCmdStr
 	}
 	shellVarMap := shellenv.ShellVarMapFromState(state)
-	if base.HasDebugFlag(shellVarMap, base.DebugFlag_LogRcFile) {
+	if base.HasDebugFlag(shellVarMap, base.DebugFlag_LogRcFile) || ForceDebugRcFile {
 		debugRcFileName := base.GetDebugRcFileName()
 		err := os.WriteFile(debugRcFileName, []byte(rcFileStr), 0600)
 		if err != nil {