code review followups: mergeEnv correctness + lifecycle and packaging fixes (#3)

* fix(agent): mergeEnv must dedup so caller overrides reach the child

The previous append(host, caller...) layout was correct in intent but wrong
under libc semantics: execve hands envp through verbatim, and getenv returns
the FIRST matching entry. With host appended before caller, any key that
existed in the host env (PATH, HOME, USER, ...) would silently shadow the
caller's override.

Switch to a map-based merge so collisions resolve to a single entry, and add
a regression test that pre-populates COCOON_AGENT_OVERRIDE_VAR in the host env
then asserts the caller value reaches the child.

* fix(cmd): silence cobra error+usage dump so exit-code failures stay quiet

run() already handles real errors (logger.Error) and translates exitCodeError
into an exit status. With cobra's defaults the latter path also prints
"Error: exit code N\nUsage:..." to stderr, which is noise — the caller already
knows the exit code and didn't ask for a usage banner.

Set SilenceErrors and SilenceUsage on the root command so only run()'s
intentional outputs reach the user.

* fix(client): prefer ctx.Err over EOF / "closed before exit" on cancel

The runCancel goroutine closes conn on ctx-cancel, which surfaces in the
readLoop's Decode as io.EOF and falls through to "agent: connection closed
before exit frame" — masking the real cause from the caller.

Gate both the in-loop EOF branch and the post-loop !sawExit branch on
ctx.Err first so caller cancel / parent timeout propagates as itself.

* fix(agent): reap Serve's ctx watcher on permanent Accept error

The previous watcher goroutine only exited via <-ctx.Done(), so a permanent
Accept failure (e.g. EMFILE, syscall-level) returned from Serve while the
goroutine kept sitting on ctx until the caller cancelled — potentially never.

Rewrite the watcher as a select on ctx.Done() vs a defer-closed done channel
so every Serve return path reaps it. Also drain connWG on the permanent-error
exit so in-flight sessions finish before the function returns.

Add TestServerWatcherExitsOnPermanentAcceptError, which uses a synthetic
listener whose Accept returns a non-net.ErrClosed error, holds the parent
ctx alive, and asserts the goroutine count returns to baseline after Serve
exits. Without the fix the test reports a leak.

* fix(cmd): log non-host vsock peer rejections + add isHostPeer test

A rejected peer used to disappear silently in the Accept loop, leaving an
operator with no signal that a misconfiguration or a guest-local probe was
hitting the listener. Both Linux and Windows listeners now log a Warn on
rejection with the peer CID/port.

Threading the logger required threading a context too — the listener owns the
Accept loop, which has no per-call ctx, so we capture the serve ctx at
listenVsock(ctx, port) construction time and store it on the struct. This is
the same ctx that the agent.Server uses, so cancellation propagates uniformly.

Add cmd/transport_linux_test.go locking in the isHostPeer contract: vsock.Host
accepted, vsock.Local rejected, non-vsock RemoteAddr rejected.

* packaging(systemd): drop wrong modprobe target + cap restart churn

vhost_vsock is the HOST-side module (driver for /dev/vhost-vsock that the
hypervisor opens), not the guest-side transport — loading it inside the guest
either fails outright or is a no-op. The guest needs virtio_transport_common
+ vmw_vsock_virtio_transport, both of which auto-load on virtio-vsock device
probe. listenVsockWithRetry already covers the bind window if the device
shows up late, so the ExecStartPre is purely misleading; drop it.

Also add StartLimitBurst=5 + StartLimitIntervalSec=60s under [Unit] (their
modern home since systemd v229) so a wedged hypervisor or viosock state
can't peg journald with restart churn — five attempts in a minute is a
firm enough boundary to surface real failures.

* test(agent): drop dead nolint:mnd, add stdin-close + framedWriter coverage

mnd is not in our enabled linter list (see .golangci.yml), so the directive
on context.WithTimeout(..., 10*time.Second) is dead weight — drop it.

Add two regression tests that were missing from prior coverage:

- TestServerMsgStdinCloseTerminatesChildStdin: spawn `wc -c`, push a known
  payload, send MsgStdinClose, assert exit=0 and the byte count came back
  intact. Locks in the mid-stream stdin-close path that wasn't exercised by
  the existing cat round-trip.

- TestFramedWriterAfterTerminal: directly drive framedWriter.Write after
  the encoder has emitted a terminal frame. The errTerminalFrameSent branch
  must NOT poison lastErr or fire cancel — otherwise the post-Wait
  err()-join would mask the legitimate exit path.

* test(agent): trim errorAcceptListener to fields the watcher test exercises

The closeMu/closed/addr fields were stub state that no caller ever read —
remove them so the fake stays focused on its single responsibility (returning
a permanent Accept error). Drops three fields, a Close mutex pair, and an
Addr override branch.

* fix(agent): tear down conns on permanent Accept error before joining

connWG.Wait on the permanent-error path could pin Serve indefinitely if a
handleConn was wedged in framedWriter.Write against a slow peer — the
ctx-cancel and net.ErrClosed paths get the listener+conn teardown for
free (watcher goroutine or external Close), but the permanent-error
branch had no such trigger.

* chore: trim verbose comments across the review branch

Most of the comments added in this branch restated the code or carried
multi-paragraph rationale that belongs in commit messages, not source.
Collapse each to a single WHY line where one is warranted, drop the
rest. -43 net comment lines, no behavior change.

* test(agent): drop t.Parallel on goroutine-leak test

NumGoroutine baseline is perturbed by sibling parallel tests in the
package (dialTestServer spawns Serve goroutines), so the <= before
assertion can stay false even with zero leak.

* test(cmd): make staticAddrConn obey io.Reader/Writer contracts

(0, nil) from Read violates io.Reader and can cause tight loops if the
stub gets reused beyond RemoteAddr(). Return io.EOF + sane Write sink.

* fix(agent): mergeEnv host dedup keeps first occurrence

Comment claimed libc-getenv semantics (first match wins) but the impl
overwrote merged[k] on every host duplicate, ending up with last-wins.
Skip if already seen so behavior matches the documented contract.

* test(agent): make watcher leak regression use specific goroutine signal

Counting all goroutines is too broad — sibling tests can perturb the
baseline. Match on the Serve watcher's stack frame instead so the test
only fires on the goroutine it actually cares about. const block lives
at the top of the file per Cocoon style.

* test(agent): bound watcher stack dump helper

runtime.Stack with an unbounded grow loop could hand back a huge buffer
on a wedged process. Cap doubling at 16 MiB so the helper still produces
a useful diagnostic without unbounded allocation.

* test(agent): simplify watcher leak polling helper

Collapse the dump-and-count helper now that callers only need the count;
the dump path is reserved for the failure diagnostic.

* build: pin golangci-lint installer to the same release tag

master's install.sh recently regressed on resolving v2.9.0, breaking
CI lint. Pinning the installer URL to the version tag matches
upstream's recommended pattern and avoids future master breakage.

Author: CMGS

Makefile

@@ -44,7 +44,7 @@ GOOSES ?= linux darwin windows
 .PHONY: golangci-lint
 golangci-lint: $(GOLANGCILINT)
 $(GOLANGCILINT):
-	curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(GOLANGCILINT_ROOT) $(GOLANGCILINT_VERSION)
+	curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/$(GOLANGCILINT_VERSION)/install.sh | sh -s -- -b $(GOLANGCILINT_ROOT) $(GOLANGCILINT_VERSION)
 
 .PHONY: gofumpt
 gofumpt: $(GOFMT)

agent/agent.go

@@ -41,10 +41,16 @@ func (s *Server) Serve(ctx context.Context) error {
 	logger := log.WithFunc("agent.Server.Serve")
 	logger.Infof(ctx, "agent listening on %s", s.listener.Addr())
 
+	// done reaps the watcher on every return path, not just ctx-cancel.
+	done := make(chan struct{})
+	defer close(done)
 	go func() {
-		<-ctx.Done()
-		_ = s.listener.Close()
-		s.closeAllConns()
+		select {
+		case <-ctx.Done():
+			_ = s.listener.Close()
+			s.closeAllConns()
+		case <-done:
+		}
 	}()
 
 	var connWG sync.WaitGroup
@@ -56,6 +62,10 @@ func (s *Server) Serve(ctx context.Context) error {
 				return nil
 			}
 			logger.Error(ctx, err, "accept")
+			// Unwedge handlers stuck on slow peers before joining.
+			_ = s.listener.Close()
+			s.closeAllConns()
+			connWG.Wait()
 			return fmt.Errorf("accept: %w", err)
 		}
 		connWG.Go(func() { s.handleConn(ctx, conn) })

agent/agent_test.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
@@ -15,6 +16,11 @@ import (
 	"github.com/cocoonstack/cocoon-agent/client"
 )
 
+const (
+	initialGoroutineDumpSize = 1 << 16
+	maxGoroutineDumpSize     = 1 << 24
+)
+
 // dialTestServer runs the agent over loopback TCP and dials a client conn.
 // Cleanup (cancel ctx, close server, wait for Serve to return, close conn)
 // is registered via t.Cleanup so callers don't repeat it per test.
@@ -25,7 +31,7 @@ func dialTestServer(t *testing.T) (context.Context, net.Conn) {
 		t.Fatalf("listen tcp: %v", err)
 	}
 	srv := agent.NewServer(tcp)
-	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) //nolint:mnd
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
 	var wg sync.WaitGroup
 	wg.Go(func() { _ = srv.Serve(ctx) })
 	conn, err := net.Dial("tcp", tcp.Addr().String())
@@ -94,6 +100,32 @@ func TestServerStreamsStdin(t *testing.T) {
 	}
 }
 
+// TestServerMsgStdinCloseTerminatesChildStdin: child must see EOF after the
+// close frame and exit 0; wc -c also confirms pre-close payload arrived.
+func TestServerMsgStdinCloseTerminatesChildStdin(t *testing.T) {
+	t.Parallel()
+	ctx, conn := dialTestServer(t)
+
+	payload := "abcde"
+	var stdout bytes.Buffer
+	exit, err := client.Run(
+		ctx, conn,
+		[]string{"sh", "-c", "wc -c"},
+		nil,
+		strings.NewReader(payload),
+		&stdout, io.Discard,
+	)
+	if err != nil {
+		t.Fatalf("client run: %v", err)
+	}
+	if exit != 0 {
+		t.Errorf("exit = %d, want 0", exit)
+	}
+	if got := strings.TrimSpace(stdout.String()); got != "5" {
+		t.Errorf("wc -c stdout = %q, want \"5\"", got)
+	}
+}
+
 func TestServerRejectsNonExecFirstFrame(t *testing.T) {
 	t.Parallel()
 	_, conn := dialTestServer(t)
@@ -224,3 +256,96 @@ func TestServerMergesEnvWithHost(t *testing.T) {
 		t.Errorf("host PATH not preserved on merge: %q", out)
 	}
 }
+
+// Not parallel: t.Setenv mutates process-global os.Environ.
+func TestServerMergesEnvCallerWins(t *testing.T) {
+	t.Setenv("COCOON_AGENT_OVERRIDE_VAR", "host-value")
+	ctx, conn := dialTestServer(t)
+
+	var stdout bytes.Buffer
+	exit, err := client.Run(
+		ctx, conn,
+		[]string{"sh", "-c", "printf %s \"$COCOON_AGENT_OVERRIDE_VAR\""},
+		map[string]string{"COCOON_AGENT_OVERRIDE_VAR": "caller-value"},
+		nil, &stdout, io.Discard,
+	)
+	if err != nil {
+		t.Fatalf("client run: %v", err)
+	}
+	if exit != 0 {
+		t.Errorf("exit = %d, want 0", exit)
+	}
+	if got := stdout.String(); got != "caller-value" {
+		t.Errorf("caller env did not win on collision: got %q, want %q", got, "caller-value")
+	}
+}
+
+// errorAcceptListener returns err on every Accept — drives Serve's
+// permanent-error return path.
+type errorAcceptListener struct {
+	err error
+}
+
+func (l *errorAcceptListener) Accept() (net.Conn, error) {
+	return nil, l.err
+}
+
+func (l *errorAcceptListener) Close() error { return nil }
+func (l *errorAcceptListener) Addr() net.Addr {
+	return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
+}
+
+func goroutineDump() string {
+	size := initialGoroutineDumpSize
+	for {
+		buf := make([]byte, size)
+		n := runtime.Stack(buf, true)
+		if n < len(buf) || size == maxGoroutineDumpSize {
+			return string(buf[:n])
+		}
+		size = min(size*2, maxGoroutineDumpSize)
+	}
+}
+
+func countServeWatcherGoroutines() int {
+	return strings.Count(goroutineDump(), "(*Server).Serve.func1")
+}
+
+// TestServerWatcherExitsOnPermanentAcceptError: on a non-ErrClosed Accept
+// failure, Serve must reap the ctx watcher via the done channel rather than
+// leak it until the (possibly never-canceled) parent ctx fires.
+//
+// Not parallel: asserts against the specific Serve watcher goroutine and keeps
+// other tests from creating extra matching stacks while it samples.
+func TestServerWatcherExitsOnPermanentAcceptError(t *testing.T) {
+	before := countServeWatcherGoroutines()
+
+	srv := agent.NewServer(&errorAcceptListener{err: errors.New("synthetic permanent accept failure")})
+	// Long-lived ctx so only the done-channel path can release the watcher.
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+
+	errCh := make(chan error, 1)
+	go func() { errCh <- srv.Serve(ctx) }()
+
+	select {
+	case err := <-errCh:
+		if err == nil {
+			t.Fatal("expected permanent accept error to surface")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Serve did not return after permanent accept error")
+	}
+
+	// Watcher exits async after close(done); poll until the specific watcher
+	// stack count returns to baseline.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		if countServeWatcherGoroutines() <= before {
+			return
+		}
+		runtime.Gosched()
+		time.Sleep(10 * time.Millisecond)
+	}
+	t.Fatalf("Serve watcher goroutine still present:\n%s", goroutineDump())
+}

agent/exec.go

@@ -5,8 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"maps"
 	"os"
 	"os/exec"
+	"strings"
 )
 
 // processController hooks platform-specific child-lifecycle steps into
@@ -144,12 +146,24 @@ func pumpStdin(ctx context.Context, w io.WriteCloser, frames <-chan Message, don
 	}
 }
 
-// mergeEnv layers caller env over os.Environ; caller keys win on collision.
-func mergeEnv(env map[string]string) []string {
-	hostEnv := os.Environ()
-	out := make([]string, 0, len(hostEnv)+len(env))
-	out = append(out, hostEnv...)
-	for k, v := range env {
+// mergeEnv layers caller env over os.Environ. Dedup is required: libc getenv
+// returns the first match, so duplicate keys would shadow caller overrides.
+func mergeEnv(callerEnv map[string]string) []string {
+	host := os.Environ()
+	merged := make(map[string]string, len(host)+len(callerEnv))
+	for _, kv := range host {
+		k, v, ok := strings.Cut(kv, "=")
+		if !ok {
+			continue
+		}
+		if _, dup := merged[k]; dup {
+			continue // first occurrence wins, matching libc getenv
+		}
+		merged[k] = v
+	}
+	maps.Copy(merged, callerEnv)
+	out := make([]string, 0, len(merged))
+	for k, v := range merged {
 		out = append(out, k+"="+v)
 	}
 	return out

agent/protocol_test.go

@@ -147,3 +147,33 @@ func TestDecodeHandlesLargeFrame(t *testing.T) {
 		t.Errorf("payload mismatch: got %d bytes, want %d", len(got.Data), len(payload))
 	}
 }
+
+// Post-terminal Write must return errTerminalFrameSent without poisoning
+// lastErr or tripping cancel, else err()-join masks the real exit.
+func TestFramedWriterAfterTerminal(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	if err := enc.Encode(Message{Type: MsgError, Message: "kaboom"}); err != nil {
+		t.Fatalf("encode terminal: %v", err)
+	}
+
+	var cancelCalled bool
+	cancel := func() { cancelCalled = true }
+	w := newFramedWriter(MsgStdout, enc, cancel)
+
+	n, err := w.Write([]byte("late stdout chunk"))
+	if !errors.Is(err, errTerminalFrameSent) {
+		t.Fatalf("post-terminal Write err = %v, want %v", err, errTerminalFrameSent)
+	}
+	if n != 0 {
+		t.Errorf("post-terminal Write n = %d, want 0", n)
+	}
+	if cancelCalled {
+		t.Error("post-terminal Write must not call cancel")
+	}
+	if got := w.err(); got != nil {
+		t.Errorf("post-terminal Write must not poison lastErr, got %v", got)
+	}
+}

client/client.go

@@ -69,12 +69,14 @@ readLoop:
 			if e := stdinReadErr.Load(); e != nil {
 				return 0, fmt.Errorf("read stdin: %w", *e)
 			}
-			if errors.Is(err, io.EOF) {
-				break
-			}
+			// Prefer ctx.Err over EOF: ctx-cancel closes the conn,
+			// surfacing as EOF here.
 			if ctx.Err() != nil {
 				return 0, ctx.Err()
 			}
+			if errors.Is(err, io.EOF) {
+				break
+			}
 			return 0, fmt.Errorf("read frame: %w", err)
 		}
 		switch frame.Type {
@@ -105,6 +107,10 @@ readLoop:
 		return 0, fmt.Errorf("read stdin: %w", *e)
 	}
 	if !sawExit {
+		// Same ctx-cancel-races-MsgExit case as the readLoop EOF path.
+		if ctx.Err() != nil {
+			return 0, ctx.Err()
+		}
 		return 0, errors.New("agent: connection closed before exit frame")
 	}
 	return exitCode, nil

cmd/root.go

@@ -21,6 +21,10 @@ func NewRootCmd() *cobra.Command {
 		Use:     "cocoon-agent",
 		Short:   "vsock-based command exec agent for Cocoon-managed VMs",
 		Version: fmt.Sprintf("%s (rev=%s built=%s)", version.VERSION, version.REVISION, version.BUILTAT),
+		// run() handles its own logging and exit codes; suppress cobra's
+		// Error/Usage dump so child-exit failures stay quiet.
+		SilenceErrors: true,
+		SilenceUsage:  true,
 	}
 	rootCmd.AddCommand(newServeCmd())
 	rootCmd.AddCommand(newClientCmd())

cmd/serve.go

@@ -44,7 +44,7 @@ func listenVsockWithRetry(ctx context.Context, port uint32) (net.Listener, error
 	logger := log.WithFunc("cmd.serve.listenRetry")
 	deadline := time.Now().Add(listenRetryTimeout)
 	for attempt := 1; ; attempt++ {
-		lsn, err := listenVsock(port)
+		lsn, err := listenVsock(ctx, port)
 		if err == nil {
 			if attempt > 1 {
 				logger.Infof(ctx, "vsock listen succeeded on attempt %d", attempt)

cmd/transport_linux.go

@@ -3,21 +3,27 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
 
 	"github.com/mdlayher/vsock"
+	"github.com/projecteru2/core/log"
 )
 
 var _ net.Listener = (*hostOnlyListener)(nil)
 
-func listenVsock(port uint32) (net.Listener, error) {
+func listenVsock(ctx context.Context, port uint32) (net.Listener, error) {
 	l, err := vsock.Listen(port, nil)
 	if err != nil {
 		return nil, fmt.Errorf("vsock listen: %w", err)
 	}
-	return &hostOnlyListener{Listener: l}, nil
+	return &hostOnlyListener{
+		Listener: l,
+		ctx:      ctx,
+		logger:   log.WithFunc("cmd.hostOnlyListener.Accept"),
+	}, nil
 }
 
 func dialVsock(cid, port uint32) (io.ReadWriteCloser, error) {
@@ -34,6 +40,9 @@ func dialVsock(cid, port uint32) (io.ReadWriteCloser, error) {
 // and trigger root-level command execution.
 type hostOnlyListener struct {
 	net.Listener
+	// ctx is the serve ctx, stashed for Accept-loop diagnostic logging.
+	ctx    context.Context
+	logger *log.Fields
 }
 
 func (l *hostOnlyListener) Accept() (net.Conn, error) {
@@ -45,6 +54,7 @@ func (l *hostOnlyListener) Accept() (net.Conn, error) {
 		if isHostPeer(conn) {
 			return conn, nil
 		}
+		l.logger.Warnf(l.ctx, "rejecting non-host vsock peer %s", conn.RemoteAddr())
 		_ = conn.Close()
 	}
 }