container/initdaemon: copy wstatus from wait4 loop

Due to the special nature of the init process, direct use of wait outside the wait4 loop is racy. This change copies the wstatus from wait4 loop state instead.

Signed-off-by: Ophestra <cat@gensokyo.uk>

4 files changed, 71 insertions, 35 deletions

diff --git a/container/dispatcher.go b/container/dispatcher.go
index e58346e8..e1dae6eb 100644
--- a/container/dispatcher.go
+++ b/container/dispatcher.go
@@ -68,8 +68,6 @@ type syscallDispatcher interface {
 	notify(c chan<- os.Signal, sig ...os.Signal)
 	// start starts [os/exec.Cmd].
 	start(c *exec.Cmd) error
-	// wait waits on [os/exec.Cmd].
-	wait(c *exec.Cmd) error
 	// signal signals the underlying process of [os/exec.Cmd].
 	signal(c *exec.Cmd, sig os.Signal) error
 	// evalSymlinks provides [filepath.EvalSymlinks].
@@ -172,7 +170,6 @@ func (direct) seccompLoad(rules []std.NativeRule, flags seccomp.ExportFlag) erro
 }
 func (direct) notify(c chan<- os.Signal, sig ...os.Signal) { signal.Notify(c, sig...) }
 func (direct) start(c *exec.Cmd) error                     { return c.Start() }
-func (direct) wait(c *exec.Cmd) error                      { return c.Wait() }
 func (direct) signal(c *exec.Cmd, sig os.Signal) error     { return c.Process.Signal(sig) }
 func (direct) evalSymlinks(path string) (string, error)    { return filepath.EvalSymlinks(path) }
 
diff --git a/container/dispatcher_test.go b/container/dispatcher_test.go
index bbcb26ad..834c700d 100644
--- a/container/dispatcher_test.go
+++ b/container/dispatcher_test.go
@@ -493,21 +493,6 @@ func (k *kstub) start(c *exec.Cmd) error {
 	return err
 }
 
-func (k *kstub) wait(c *exec.Cmd) error {
-	k.Helper()
-	expect := k.Expects("wait")
-	err := expect.Error(
-		stub.CheckArg(k.Stub, "c.Path", c.Path, 0),
-		stub.CheckArgReflect(k.Stub, "c.Args", c.Args, 1),
-		stub.CheckArgReflect(k.Stub, "c.Env", c.Env, 2),
-		stub.CheckArg(k.Stub, "c.Dir", c.Dir, 3))
-
-	if mgc, ok := expect.Ret.(uintptr); ok && mgc == stub.PanicExit {
-		panic(stub.PanicExit)
-	}
-	return err
-}
-
 func (k *kstub) signal(c *exec.Cmd, sig os.Signal) error {
 	k.Helper()
 	expect := k.Expects("signal")
diff --git a/container/initdaemon.go b/container/initdaemon.go
index ee360dc9..0e631b25 100644
--- a/container/initdaemon.go
+++ b/container/initdaemon.go
@@ -8,7 +8,7 @@ import (
 	"os"
 	"os/exec"
 	"slices"
-	"sync/atomic"
+	"strconv"
 	"syscall"
 	"time"
 
@@ -41,6 +41,38 @@ type DaemonOp struct {
 	Args []string
 }
 
+// earlyTerminationError is returned by [DaemonOp] when a daemon terminates
+// before [DaemonOp.Target] appears.
+type earlyTerminationError struct {
+	// Returned by [DaemonOp.String].
+	op string
+	// Copied from wait4 loop.
+	wstatus syscall.WaitStatus
+}
+
+func (e *earlyTerminationError) Error() string {
+	res := ""
+	switch {
+	case e.wstatus.Exited():
+		res = "exit status " + strconv.Itoa(e.wstatus.ExitStatus())
+	case e.wstatus.Signaled():
+		res = "signal: " + e.wstatus.Signal().String()
+	case e.wstatus.Stopped():
+		res = "stop signal: " + e.wstatus.StopSignal().String()
+		if e.wstatus.StopSignal() == syscall.SIGTRAP && e.wstatus.TrapCause() != 0 {
+			res += " (trap " + strconv.Itoa(e.wstatus.TrapCause()) + ")"
+		}
+	case e.wstatus.Continued():
+		res = "continued"
+	}
+	if e.wstatus.CoreDump() {
+		res += " (core dumped)"
+	}
+	return res
+}
+
+func (e *earlyTerminationError) Message() string { return e.op + " " + e.Error() }
+
 func (d *DaemonOp) Valid() bool                                { return d != nil && d.Target != nil && d.Path != nil }
 func (d *DaemonOp) early(*setupState, syscallDispatcher) error { return nil }
 func (d *DaemonOp) apply(*setupState, syscallDispatcher) error { return nil }
@@ -59,17 +91,9 @@ func (d *DaemonOp) late(state *setupState, k syscallDispatcher) error {
 		return err
 	}
 
-	var done atomic.Pointer[error]
-	k.new(func(k syscallDispatcher) {
-		err := k.wait(cmd)
-		done.Store(&err)
-
-		if err != nil {
-			state.Verbosef("%s %v", d.String(), err)
-		}
-	})
-
 	deadline := time.Now().Add(daemonTimeout)
+	var wstatusErr error
+
 	for {
 		if _, err := k.stat(d.Target.String()); err != nil {
 			if !errors.Is(err, os.ErrNotExist) {
@@ -82,8 +106,13 @@ func (d *DaemonOp) late(state *setupState, k syscallDispatcher) error {
 				return context.DeadlineExceeded
 			}
 
-			if errP := done.Load(); errP != nil {
-				return *errP
+			if wstatusErr != nil {
+				return wstatusErr
+			}
+			if wstatus, ok := state.terminated(cmd.Process.Pid); ok {
+				// check once again: process could have satisfied Target between stat and the lookup
+				wstatusErr = &earlyTerminationError{d.String(), wstatus}
+				continue
 			}
 
 			time.Sleep(500 * time.Microsecond)
diff --git a/container/initdaemon_test.go b/container/initdaemon_test.go
index 5ab834db..ad8c7a93 100644
--- a/container/initdaemon_test.go
+++ b/container/initdaemon_test.go
@@ -6,8 +6,36 @@ import (
 
 	"hakurei.app/container/check"
 	"hakurei.app/container/stub"
+	"hakurei.app/message"
 )
 
+func TestEarlyTerminationError(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name string
+		err  error
+		want string
+		msg  string
+	}{
+		{"exited", &earlyTerminationError{
+			`daemon providing "/run/user/1971/pulse/native"`, 127 << 8,
+		}, "exit status 127", `daemon providing "/run/user/1971/pulse/native" exit status 127`},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			if got := tc.err.Error(); got != tc.want {
+				t.Errorf("Error: %q, want %q", got, tc.want)
+			}
+			if got := tc.err.(message.Error).Message(); got != tc.msg {
+				t.Errorf("Message: %s, want %s", got, tc.msg)
+			}
+		})
+	}
+}
+
 func TestDaemonOp(t *testing.T) {
 	t.Parallel()
 
@@ -23,15 +51,12 @@ func TestDaemonOp(t *testing.T) {
 			call("isVerbose", stub.ExpectArgs{}, true, nil),
 			call("verbosef", stub.ExpectArgs{"starting %s", []any{`daemon providing "/run/user/1971/pulse/native"`}}, nil, nil),
 			call("start", stub.ExpectArgs{"/run/current-system/sw/bin/pipewire-pulse", []string{"/run/current-system/sw/bin/pipewire-pulse", "-v"}, []string{"\x00"}, "/"}, &os.Process{Pid: 0xcafe}, nil),
-			call("New", stub.ExpectArgs{}, nil, nil),
 			call("stat", stub.ExpectArgs{"/run/user/1971/pulse/native"}, isDirFi(false), os.ErrNotExist),
 			call("stat", stub.ExpectArgs{"/run/user/1971/pulse/native"}, isDirFi(false), os.ErrNotExist),
 			call("stat", stub.ExpectArgs{"/run/user/1971/pulse/native"}, isDirFi(false), os.ErrNotExist),
 			call("stat", stub.ExpectArgs{"/run/user/1971/pulse/native"}, isDirFi(false), nil),
 			call("verbosef", stub.ExpectArgs{"daemon process %d ready", []any{0xcafe}}, nil, nil),
-		}, Tracks: []stub.Expect{{Calls: []stub.Call{
-			call("wait", stub.ExpectArgs{"/run/current-system/sw/bin/pipewire-pulse", []string{"/run/current-system/sw/bin/pipewire-pulse", "-v"}, []string{"\x00"}, "/"}, uintptr(stub.PanicExit), nil),
-		}}}}, nil},
+		}}, nil},
 	})
 
 	checkOpsValid(t, []opValidTestCase{