fix(serve): forward the full environment to the streaming daemon

src/cli/serve.go

@@ -26,10 +26,15 @@ func init() {
 	RootCmd.AddCommand(serveCmd)
 }
 
-// One JSON object per line on stdin. Unknown fields are ignored by
-// encoding/json by default, giving forward compatibility for free.
+// One JSON object per line on stdin, immediately followed - for every
+// command, not just render - by a raw "KEY=VALUE\x00" record stream
+// terminated by an empty record (a bare NUL); see readEnvBlob. Unknown JSON
+// fields are ignored by encoding/json by default, giving forward
+// compatibility for free.
 type serveRequest struct {
-	Env           map[string]string `json:"env"`
+	// Env is never part of the JSON header - it comes from the raw record
+	// stream that follows every request line, parsed by readEnvBlob.
+	Env           map[string]string `json:"-"`
 	Command       string            `json:"command"`
 	Shell         string            `json:"shell"`
 	ShellVersion  string            `json:"shell-version"`
@@ -112,9 +117,11 @@ func createServeCmd() *cmdtree.Command {
 // primitive fish has) never EOFs the read side. Unix only - the shell owns
 // the fifo's lifecycle.
 //
-// Clients must write each request in a single write(2) call; requests from a
-// single sequential writer (one shell session) never interleave regardless
-// of size.
+// A request (header line plus its env blob, see readEnvBlob) may span more
+// than one write(2) call - the reader is not line/buffer-size bound - but
+// those calls must be consecutive with no other writer's bytes landing
+// between them. A single sequential writer (one shell session, one request
+// at a time) guarantees that regardless of size.
 func openServeInput(pipePath string) (*os.File, error) {
 	if pipePath == "" {
 		return os.Stdin, nil
@@ -139,7 +146,8 @@ type serveActiveCycle struct {
 	copierDone chan struct{}
 }
 
-// runServeLoop reads newline-delimited JSON requests from in and writes
+// runServeLoop reads newline-delimited JSON requests (each immediately
+// followed by a raw env record blob, see readEnvBlob) from in and writes
 // NUL-delimited, cycle-id-prefixed prompt records to out. It returns when it
 // reads a quit command or hits EOF on stdin. The returned bool reports
 // whether at least one render request was handled, so the caller knows
@@ -152,20 +160,18 @@ type serveActiveCycle struct {
 // shell additionally redirects this process's stderr so anything unrecovered
 // can never reach the user's terminal.
 func runServeLoop(in, out *os.File) bool {
-	scanner := bufio.NewScanner(in)
-	// Env payloads (a POSH_* overlay plus PATH) can exceed the default 64 KB
-	// scanner buffer, so grow it up front.
-	scanner.Buffer(make([]byte, 0, 256*1024), 1024*1024)
+	reader := bufio.NewReader(in)
 
 	var active *serveActiveCycle
 	renderedAtLeastOnce := false
 
-	// envKeys tracks which variables the previous request's overlay set, so
+	// envKeys tracks which variables the previous request's env blob set, so
 	// a variable that disappears from a later request (e.g. VIRTUAL_ENV after
-	// `deactivate`) gets unset instead of pinning its stale value for the rest
-	// of the daemon's life. Scoped to the loop so repeated invocations in the
-	// same process (tests) never inherit a previous loop's keys. The serve
-	// loop is single-threaded, so no locking.
+	// `deactivate`, or anything a client stops forwarding) gets unset instead
+	// of pinning its stale value for the rest of the daemon's life. Scoped to
+	// the loop so repeated invocations in the same process (tests) never
+	// inherit a previous loop's keys. The serve loop is single-threaded, so
+	// no locking.
 	envKeys := map[string]struct{}{}
 
 	stopActiveCycle := func() {
@@ -189,52 +195,108 @@ func runServeLoop(in, out *os.File) bool {
 		active = nil
 	}
 
-	for scanner.Scan() {
-		line := scanner.Bytes()
+	for {
+		line, err := reader.ReadBytes('\n')
+		eof := err != nil
 
+		line = bytes.TrimSuffix(line, []byte{'\n'})
+		line = bytes.TrimSuffix(line, []byte{'\r'})
 		// Strip a UTF-8 BOM: .NET's default UTF8 encoding writes one on the
 		// StreamWriter's first write, which would otherwise make the first
 		// request line of a session unparseable.
 		line = bytes.TrimPrefix(line, []byte{0xEF, 0xBB, 0xBF})
 
 		if len(line) == 0 {
+			if eof {
+				break
+			}
 			continue
 		}
 
-		var req serveRequest
-		if err := json.Unmarshal(line, &req); err != nil {
-			// Malformed line: ignore for forward/backward compatibility.
-			continue
+		// A well-formed client always sends the env blob right after the
+		// header line, for every command - even abort/quit send a bare NUL
+		// terminator. Reading it here, unconditionally, is what keeps the
+		// stream in sync regardless of the header's command or JSON validity;
+		// a client that skipped it on some commands would desync every
+		// request after the first one that did.
+		env, envErr := readEnvBlob(reader)
+		if envErr != nil {
+			// Truncated/closed mid-blob: nothing more can be recovered.
+			break
 		}
 
-		switch req.Command {
-		case serveCommandRender:
-			// A new render request implicitly aborts whatever is running.
-			stopActiveCycle()
-			// A nil cycle means setup panicked before prompt.New completed -
-			// template.Init may never have run, in which case the shutdown
-			// path must not call template.SaveCache (it dereferences state
-			// only Init sets). A started cycle implies Init completed.
-			if active = startRenderCycle(&req, out, envKeys); active != nil {
-				renderedAtLeastOnce = true
+		var req serveRequest
+		if err := json.Unmarshal(line, &req); err == nil {
+			req.Env = env
+
+			switch req.Command {
+			case serveCommandRender:
+				// A new render request implicitly aborts whatever is running.
+				stopActiveCycle()
+				// A nil cycle means setup panicked before prompt.New completed -
+				// template.Init may never have run, in which case the shutdown
+				// path must not call template.SaveCache (it dereferences state
+				// only Init sets). A started cycle implies Init completed.
+				if active = startRenderCycle(&req, out, envKeys); active != nil {
+					renderedAtLeastOnce = true
+				}
+			case serveCommandAbort:
+				stopActiveCycle()
+			case serveCommandQuit:
+				stopActiveCycle()
+				return renderedAtLeastOnce
+			default:
+				// Unknown command: ignore for forward compatibility.
 			}
-		case serveCommandAbort:
-			stopActiveCycle()
-		case serveCommandQuit:
-			stopActiveCycle()
-			return renderedAtLeastOnce
-		default:
-			// Unknown command: ignore for forward compatibility.
+		}
+		// Malformed JSON header: ignored for forward/backward compatibility
+		// (its env blob was already consumed above, keeping the stream in sync).
+
+		if eof {
+			break
 		}
 	}
 
-	// EOF (or a scanner error) on stdin: behave like an explicit quit so
-	// caches are still flushed by the caller's deferred cleanup.
+	// EOF (or a read error) on stdin: behave like an explicit quit so caches
+	// are still flushed by the caller's deferred cleanup.
 	stopActiveCycle()
 
 	return renderedAtLeastOnce
 }
 
+// readEnvBlob reads a "KEY=VALUE\x00" record stream from r, terminated by an
+// empty record (a bare NUL byte). Every request line is unconditionally
+// followed by this blob - even for commands that ignore its contents - so
+// the reader never needs to know in advance whether one is coming.
+//
+// Environment variable values cannot contain a NUL byte on any OS this
+// project targets (POSIX environ entries and the Windows environment block
+// are themselves NUL-terminated/-delimited C strings), so this framing needs
+// no escaping: a key/value pair is malformed only if it has no '=', in which
+// case it is skipped.
+func readEnvBlob(r *bufio.Reader) (map[string]string, error) {
+	env := map[string]string{}
+
+	for {
+		record, err := r.ReadBytes(0)
+		if err != nil {
+			return nil, err
+		}
+
+		record = record[:len(record)-1] // drop the trailing NUL delimiter
+		if len(record) == 0 {
+			return env, nil
+		}
+
+		key, value, found := bytes.Cut(record, []byte{'='})
+		if !found {
+			continue
+		}
+
+		env[string(key)] = string(value)
+	}
+}
+
 func applyEnvOverlay(env map[string]string, keys map[string]struct{}) {
 	for key := range keys {
 		if _, ok := env[key]; ok {

src/cli/serve_pipe_test.go

@@ -53,7 +53,9 @@ func TestServeLoop_RequestPipe(t *testing.T) {
 		data, err := json.Marshal(v)
 		require.NoError(t, err)
 
-		_, err = f.Write(append(data, '\n'))
+		data = append(data, '\n', 0) // trailing 0: empty env blob, just the terminator
+
+		_, err = f.Write(data)
 		require.NoError(t, err)
 		require.NoError(t, f.Close())
 	}

src/cli/serve_test.go

@@ -59,14 +59,40 @@ func startServeHarness(t *testing.T) *serveHarness {
 	return h
 }
 
-// send writes a single newline-terminated JSON request to the loop's stdin.
+// send writes a newline-terminated JSON header to the loop's stdin, followed
+// by the NUL-delimited env blob every request must carry (see readEnvBlob).
+// v may carry an "env" key (map[string]string) - if present, it is pulled
+// out of the JSON header and sent as the raw blob instead; a request with no
+// "env" key sends an empty blob (just the terminator).
 func (h *serveHarness) send(v any) {
 	h.t.Helper()
 
+	env := map[string]string{}
+	if m, ok := v.(map[string]any); ok {
+		if raw, ok := m["env"]; ok {
+			delete(m, "env")
+			typed, ok := raw.(map[string]string)
+			require.True(h.t, ok, "send: \"env\" must be a map[string]string, got %T", raw)
+			env = typed
+		}
+	}
+
 	data, err := json.Marshal(v)
 	require.NoError(h.t, err)
 
-	_, err = h.stdin.Write(append(data, '\n'))
+	var buf bytes.Buffer
+	buf.Write(data)
+	buf.WriteByte('\n')
+
+	for key, value := range env {
+		buf.WriteString(key)
+		buf.WriteByte('=')
+		buf.WriteString(value)
+		buf.WriteByte(0)
+	}
+	buf.WriteByte(0) // empty record: terminates the blob
+
+	_, err = h.stdin.Write(buf.Bytes())
 	require.NoError(h.t, err)
 }
 
@@ -430,6 +456,91 @@ func TestServeLoop_EnvOverlayUnsetsVanishedVariables(t *testing.T) {
 	h.quitAndWait()
 }
 
+// TestServeLoop_EnvBlobHandlesArbitraryValues guards the reason env forwarding
+// moved off JSON: a value with a literal newline, tab, quote, backslash, or
+// non-ASCII byte must reach the daemon byte-exact, with no escaping logic to
+// get wrong. This would corrupt or silently drop such values under JSON
+// string escaping (which is exactly what the whitelist-based overlay hid,
+// since it only ever carried PATH/POSH_*-shaped values).
+func TestServeLoop_EnvBlobHandlesArbitraryValues(t *testing.T) {
+	h := startServeHarness(t)
+	pwd := t.TempDir()
+	chdirBackToWD(t)
+
+	const name = "POSH_SERVE_ENV_ARBITRARY_TEST"
+	t.Cleanup(func() { _ = os.Unsetenv(name) })
+
+	value := "line1\nline2\ttab\"quote\\backénd"
+
+	h.send(map[string]any{
+		"command": "render", "id": 1, "shell": "pwsh", "pwd": pwd,
+		"env": map[string]string{name: value},
+	})
+	records := h.records(500 * time.Millisecond)
+	require.NotEmpty(t, records)
+	assert.Equal(t, value, os.Getenv(name), "env value must survive the blob byte-exact, unescaped")
+
+	h.quitAndWait()
+}
+
+// TestReadEnvBlob_MalformedRecordsSkipped guards readEnvBlob's parsing rules
+// directly: a record with no '=' is dropped rather than corrupting the map
+// or aborting the parse, and an empty blob (just the terminator) parses to
+// an empty, non-nil map.
+func TestReadEnvBlob_MalformedRecordsSkipped(t *testing.T) {
+	blob := "NOEQUALSSIGN\x00KEY=value\x00ANOTHER=a=b=c\x00\x00"
+	reader := bufio.NewReader(bytes.NewBufferString(blob))
+
+	env, err := readEnvBlob(reader)
+	require.NoError(t, err)
+	assert.NotContains(t, env, "NOEQUALSSIGN", "a record with no '=' must be skipped, not stored under an empty value")
+	assert.Equal(t, "value", env["KEY"])
+	assert.Equal(t, "a=b=c", env["ANOTHER"], "only the first '=' splits key from value")
+
+	empty, err := readEnvBlob(bufio.NewReader(bytes.NewBufferString("\x00")))
+	require.NoError(t, err)
+	assert.Empty(t, empty)
+}
+
+// TestReadEnvBlob_TruncatedBlobReturnsError guards the case where the
+// connection closes mid-record, before the terminating empty record ever
+// arrives: readEnvBlob must report an error (so runServeLoop treats it like
+// EOF and shuts down) rather than block forever or return a partial map.
+func TestReadEnvBlob_TruncatedBlobReturnsError(t *testing.T) {
+	reader := bufio.NewReader(bytes.NewBufferString("KEY=value\x00TRUNC=no-terminator-ever"))
+
+	env, err := readEnvBlob(reader)
+	require.Error(t, err)
+	assert.Nil(t, env)
+}
+
+// TestServeLoop_MalformedHeaderStillConsumesBlobThenNextRequestRenders is
+// the central desync-resilience property runServeLoop's comments claim: a
+// non-empty header line that fails JSON parsing must still have its env blob
+// consumed (every header is unconditionally followed by one), or the next
+// request's header would be misread as more of the previous blob and the
+// loop would never render again.
+func TestServeLoop_MalformedHeaderStillConsumesBlobThenNextRequestRenders(t *testing.T) {
+	h := startServeHarness(t)
+	pwd := t.TempDir()
+	chdirBackToWD(t)
+
+	// A malformed (non-JSON) header line, followed by its own non-empty env
+	// blob - exactly what a well-formed client always sends, just with a
+	// garbled header. The header still needs its own newline terminator;
+	// omitting it would just make this one long header line, defeating the
+	// point of the test.
+	_, err := h.stdin.Write([]byte("not valid json\nSOME=value\x00\x00"))
+	require.NoError(t, err)
+
+	h.render(1, pwd)
+	records := h.records(500 * time.Millisecond)
+	require.NotEmpty(t, records, "a render after a malformed header must still produce records - the stream must not have desynced")
+	assert.Equal(t, "1", records[0].id)
+
+	h.quitAndWait()
+}
+
 func TestServeLoop_QuitExitsCleanly(t *testing.T) {
 	h := startServeHarness(t)
 
@@ -496,7 +607,7 @@ func TestServeLoop_UTF8BOMOnFirstLine(t *testing.T) {
 	require.NoError(t, err)
 
 	payload := append([]byte{0xEF, 0xBB, 0xBF}, data...)
-	payload = append(payload, '\n')
+	payload = append(payload, '\n', 0) // trailing 0: empty env blob, just the terminator
 	_, err = h.stdin.Write(payload)
 	require.NoError(t, err)