helper: embed context on creation

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

8 files changed, 76 insertions, 75 deletions

diff --git a/dbus/run.go b/dbus/run.go
index 074fb92d..fc3943a5 100644
--- a/dbus/run.go
+++ b/dbus/run.go
@@ -37,8 +37,9 @@ func (p *Proxy) Start(ctx context.Context, output io.Writer, sandbox bool) error
 		}
 	)
 
+	c, cancel := context.WithCancelCause(ctx)
 	if !sandbox {
-		h = helper.New(p.seal, p.name, argF)
+		h = helper.New(c, p.seal, p.name, argF)
 		// xdg-dbus-proxy does not need to inherit the environment
 		h.SetEnv(make([]string, 0))
 	} else {
@@ -110,15 +111,14 @@ func (p *Proxy) Start(ctx context.Context, output io.Writer, sandbox bool) error
 			bc.Bind(k, k)
 		}
 
-		h = helper.MustNewBwrap(bc, toolPath, true, p.seal, argF, nil, nil)
+		h = helper.MustNewBwrap(c, bc, toolPath, true, p.seal, argF, nil, nil)
 		p.bwrap = bc
 	}
 
 	if output != nil {
-		h.Stdout(output).Stderr(output)
+		h.SetStdout(output).SetStderr(output)
 	}
-	c, cancel := context.WithCancelCause(ctx)
-	if err := h.Start(c, true); err != nil {
+	if err := h.Start(true); err != nil {
 		cancel(err)
 		return err
 	}
diff --git a/helper/bwrap.go b/helper/bwrap.go
index 8ce224bd..86aff4db 100644
--- a/helper/bwrap.go
+++ b/helper/bwrap.go
@@ -31,7 +31,7 @@ type bubblewrap struct {
 	*helperCmd
 }
 
-func (b *bubblewrap) Start(ctx context.Context, stat bool) error {
+func (b *bubblewrap) Start(stat bool) error {
 	b.lock.Lock()
 	defer b.lock.Unlock()
 
@@ -41,27 +41,24 @@ func (b *bubblewrap) Start(ctx context.Context, stat bool) error {
 		return errors.New("exec: already started")
 	}
 
-	args := b.finalise(ctx, stat)
-	if b.setpgid {
-		b.Cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
-	}
-
+	args := b.finalise(stat)
 	b.Cmd.Args = slices.Grow(b.Cmd.Args, 4+len(args))
 	b.Cmd.Args = append(b.Cmd.Args, "--args", strconv.Itoa(int(b.argsFd)), "--", b.name)
 	b.Cmd.Args = append(b.Cmd.Args, args...)
-	return proc.Fulfill(ctx, b.Cmd, b.files, b.extraFiles)
+	return proc.Fulfill(b.ctx, b.Cmd, b.files, b.extraFiles)
 }
 
 // MustNewBwrap initialises a new Bwrap instance with wt as the null-terminated argument writer.
 // If wt is nil, the child process spawned by bwrap will not get an argument pipe.
 // Function argF returns an array of arguments passed directly to the child process.
 func MustNewBwrap(
+	ctx context.Context,
 	conf *bwrap.Config, name string, setpgid bool,
 	wt io.WriterTo, argF func(argsFD, statFD int) []string,
 	extraFiles []*os.File,
 	syncFd *os.File,
 ) Helper {
-	b, err := NewBwrap(conf, name, setpgid, wt, argF, extraFiles, syncFd)
+	b, err := NewBwrap(ctx, conf, name, setpgid, wt, argF, extraFiles, syncFd)
 	if err != nil {
 		panic(err.Error())
 	} else {
@@ -73,6 +70,7 @@ func MustNewBwrap(
 // If wt is nil, the child process spawned by bwrap will not get an argument pipe.
 // Function argF returns an array of arguments passed directly to the child process.
 func NewBwrap(
+	ctx context.Context,
 	conf *bwrap.Config, name string, setpgid bool,
 	wt io.WriterTo, argF func(argsFd, statFd int) []string,
 	extraFiles []*os.File,
@@ -82,7 +80,10 @@ func NewBwrap(
 
 	b.name = name
 	b.setpgid = setpgid
-	b.helperCmd = newHelperCmd(b, BubblewrapName, wt, argF, extraFiles)
+	b.helperCmd = newHelperCmd(b, ctx, BubblewrapName, wt, argF, extraFiles)
+	if b.setpgid {
+		b.Cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	}
 
 	if v, err := NewCheckedArgs(conf.Args(syncFd, b.extraFiles, &b.files)); err != nil {
 		return nil, err
diff --git a/helper/bwrap_test.go b/helper/bwrap_test.go
index 6426434a..a14a4179 100644
--- a/helper/bwrap_test.go
+++ b/helper/bwrap_test.go
@@ -31,12 +31,13 @@ func TestBwrap(t *testing.T) {
 		})
 
 		h := helper.MustNewBwrap(
+			context.Background(),
 			sc, "fortify", false,
 			argsWt, argF,
 			nil, nil,
 		)
 
-		if err := h.Start(context.Background(), false); !errors.Is(err, os.ErrNotExist) {
+		if err := h.Start(false); !errors.Is(err, os.ErrNotExist) {
 			t.Errorf("Start: error = %v, wantErr %v",
 				err, os.ErrNotExist)
 		}
@@ -44,6 +45,7 @@ func TestBwrap(t *testing.T) {
 
 	t.Run("valid new helper nil check", func(t *testing.T) {
 		if got := helper.MustNewBwrap(
+			context.TODO(),
 			sc, "fortify", false,
 			argsWt, argF,
 			nil, nil,
@@ -64,6 +66,7 @@ func TestBwrap(t *testing.T) {
 		}()
 
 		helper.MustNewBwrap(
+			context.TODO(),
 			&bwrap.Config{Hostname: "\x00"}, "fortify", false,
 			nil, argF,
 			nil, nil,
@@ -73,19 +76,19 @@ func TestBwrap(t *testing.T) {
 	t.Run("start without pipes", func(t *testing.T) {
 		helper.InternalReplaceExecCommand(t)
 
+		c, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
 		h := helper.MustNewBwrap(
+			c,
 			sc, "crash-test-dummy", false,
 			nil, argFChecked,
 			nil, nil,
 		)
 
 		stdout, stderr := new(strings.Builder), new(strings.Builder)
-		h.Stdout(stdout).Stderr(stderr)
-
-		c, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
+		h.SetStdout(stdout).SetStderr(stderr)
 
-		if err := h.Start(c, false); err != nil {
+		if err := h.Start(false); err != nil {
 			t.Errorf("Start: error = %v",
 				err)
 			return
@@ -98,8 +101,9 @@ func TestBwrap(t *testing.T) {
 	})
 
 	t.Run("implementation compliance", func(t *testing.T) {
-		testHelper(t, func() helper.Helper {
+		testHelper(t, func(ctx context.Context) helper.Helper {
 			return helper.MustNewBwrap(
+				ctx,
 				sc, "crash-test-dummy", false,
 				argsWt, argF, nil, nil,
 			)
diff --git a/helper/direct.go b/helper/direct.go
index 9198952d..fcc7b4e8 100644
--- a/helper/direct.go
+++ b/helper/direct.go
@@ -16,7 +16,7 @@ type direct struct {
 	*helperCmd
 }
 
-func (h *direct) Start(ctx context.Context, stat bool) error {
+func (h *direct) Start(stat bool) error {
 	h.lock.Lock()
 	defer h.lock.Unlock()
 
@@ -26,15 +26,15 @@ func (h *direct) Start(ctx context.Context, stat bool) error {
 		return errors.New("exec: already started")
 	}
 
-	args := h.finalise(ctx, stat)
+	args := h.finalise(stat)
 	h.Cmd.Args = append(h.Cmd.Args, args...)
-	return proc.Fulfill(ctx, h.Cmd, h.files, h.extraFiles)
+	return proc.Fulfill(h.ctx, h.Cmd, h.files, h.extraFiles)
 }
 
 // New initialises a new direct Helper instance with wt as the null-terminated argument writer.
 // Function argF returns an array of arguments passed directly to the child process.
-func New(wt io.WriterTo, name string, argF func(argsFd, statFd int) []string) Helper {
+func New(ctx context.Context, wt io.WriterTo, name string, argF func(argsFd, statFd int) []string) Helper {
 	d := new(direct)
-	d.helperCmd = newHelperCmd(d, name, wt, argF, nil)
+	d.helperCmd = newHelperCmd(d, ctx, name, wt, argF, nil)
 	return d
 }
diff --git a/helper/direct_test.go b/helper/direct_test.go
index 67b86827..e7bda797 100644
--- a/helper/direct_test.go
+++ b/helper/direct_test.go
@@ -11,16 +11,16 @@ import (
 
 func TestDirect(t *testing.T) {
 	t.Run("start non-existent helper path", func(t *testing.T) {
-		h := helper.New(argsWt, "/nonexistent", argF)
+		h := helper.New(context.Background(), argsWt, "/nonexistent", argF)
 
-		if err := h.Start(context.Background(), false); !errors.Is(err, os.ErrNotExist) {
+		if err := h.Start(false); !errors.Is(err, os.ErrNotExist) {
 			t.Errorf("Start: error = %v, wantErr %v",
 				err, os.ErrNotExist)
 		}
 	})
 
 	t.Run("valid new helper nil check", func(t *testing.T) {
-		if got := helper.New(argsWt, "fortify", argF); got == nil {
+		if got := helper.New(context.TODO(), argsWt, "fortify", argF); got == nil {
 			t.Errorf("New(%q, %q) got nil",
 				argsWt, "fortify")
 			return
@@ -28,6 +28,6 @@ func TestDirect(t *testing.T) {
 	})
 
 	t.Run("implementation compliance", func(t *testing.T) {
-		testHelper(t, func() helper.Helper { return helper.New(argsWt, "crash-test-dummy", argF) })
+		testHelper(t, func(ctx context.Context) helper.Helper { return helper.New(ctx, argsWt, "crash-test-dummy", argF) })
 	})
 }
diff --git a/helper/helper.go b/helper/helper.go
index 2ef1afaf..f065d0ef 100644
--- a/helper/helper.go
+++ b/helper/helper.go
@@ -1,4 +1,4 @@
-// Package helper runs external helpers with optional sandboxing and manages their status/args pipes.
+// Package helper runs external helpers with optional sandboxing.
 package helper
 
 import (
@@ -26,18 +26,18 @@ const (
 )
 
 type Helper interface {
-	// Stdin sets the standard input of Helper.
-	Stdin(r io.Reader) Helper
-	// Stdout sets the standard output of Helper.
-	Stdout(w io.Writer) Helper
-	// Stderr sets the standard error of Helper.
-	Stderr(w io.Writer) Helper
+	// SetStdin sets the standard input of Helper.
+	SetStdin(r io.Reader) Helper
+	// SetStdout sets the standard output of Helper.
+	SetStdout(w io.Writer) Helper
+	// SetStderr sets the standard error of Helper.
+	SetStderr(w io.Writer) Helper
 	// SetEnv sets the environment of Helper.
 	SetEnv(env []string) Helper
 
 	// Start starts the helper process.
 	// A status pipe is passed to the helper if stat is true.
-	Start(ctx context.Context, stat bool) error
+	Start(stat bool) error
 	// Wait blocks until Helper exits and releases all its resources.
 	Wait() error
 
@@ -45,14 +45,17 @@ type Helper interface {
 }
 
 func newHelperCmd(
-	h Helper, name string,
+	h Helper, ctx context.Context, name string,
 	wt io.WriterTo, argF func(argsFd, statFd int) []string,
 	extraFiles []*os.File,
 ) (cmd *helperCmd) {
 	cmd = new(helperCmd)
-
 	cmd.r = h
-	cmd.name = name
+	cmd.ctx = ctx
+
+	cmd.Cmd = commandContext(ctx, name)
+	cmd.Cmd.Cancel = func() error { return cmd.Process.Signal(syscall.SIGTERM) }
+	cmd.WaitDelay = WaitDelay
 
 	cmd.extraFiles = new(proc.ExtraFilesPre)
 	for _, f := range extraFiles {
@@ -90,32 +93,24 @@ type helperCmd struct {
 	// passed through to [proc.Fulfill] and [proc.InitFile]
 	extraFiles *proc.ExtraFilesPre
 
-	name           string
-	stdin          io.Reader
-	stdout, stderr io.Writer
-	env            []string
+	ctx context.Context
 	*exec.Cmd
 }
 
-func (h *helperCmd) Stdin(r io.Reader) Helper   { h.stdin = r; return h.r }
-func (h *helperCmd) Stdout(w io.Writer) Helper  { h.stdout = w; return h.r }
-func (h *helperCmd) Stderr(w io.Writer) Helper  { h.stderr = w; return h.r }
-func (h *helperCmd) SetEnv(env []string) Helper { h.env = env; return h.r }
+func (h *helperCmd) SetStdin(r io.Reader) Helper  { h.Stdin = r; return h.r }
+func (h *helperCmd) SetStdout(w io.Writer) Helper { h.Stdout = w; return h.r }
+func (h *helperCmd) SetStderr(w io.Writer) Helper { h.Stderr = w; return h.r }
+func (h *helperCmd) SetEnv(env []string) Helper   { h.Env = env; return h.r }
 
-// finalise initialises the underlying [exec.Cmd] object.
-func (h *helperCmd) finalise(ctx context.Context, stat bool) (args []string) {
-	h.Cmd = commandContext(ctx, h.name)
-	h.Cmd.Stdin, h.Cmd.Stdout, h.Cmd.Stderr = h.stdin, h.stdout, h.stderr
-	h.Cmd.Env = slices.Grow(h.env, 2)
+// finalise sets up the underlying [exec.Cmd] object.
+func (h *helperCmd) finalise(stat bool) (args []string) {
+	h.Env = slices.Grow(h.Env, 2)
 	if h.hasArgsFd {
-		h.Cmd.Env = append(h.Cmd.Env, FortifyHelper+"=1")
+		h.Cmd.Env = append(h.Env, FortifyHelper+"=1")
 	} else {
-		h.Cmd.Env = append(h.Cmd.Env, FortifyHelper+"=0")
+		h.Cmd.Env = append(h.Env, FortifyHelper+"=0")
 	}
 
-	h.Cmd.Cancel = func() error { return h.Cmd.Process.Signal(syscall.SIGTERM) }
-	h.Cmd.WaitDelay = WaitDelay
-
 	statFd := -1
 	if stat {
 		f := proc.NewStat(&h.stat)
diff --git a/helper/helper_test.go b/helper/helper_test.go
index 69be474e..363cc04d 100644
--- a/helper/helper_test.go
+++ b/helper/helper_test.go
@@ -46,14 +46,15 @@ func argFChecked(argsFd, statFd int) (args []string) {
 }
 
 // this function tests an implementation of the helper.Helper interface
-func testHelper(t *testing.T, createHelper func() helper.Helper) {
+func testHelper(t *testing.T, createHelper func(ctx context.Context) helper.Helper) {
 	helper.InternalReplaceExecCommand(t)
 
 	t.Run("start helper with status channel and wait", func(t *testing.T) {
-		h := createHelper()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		h := createHelper(ctx)
 
 		stdout, stderr := new(strings.Builder), new(strings.Builder)
-		h.Stdout(stdout).Stderr(stderr)
+		h.SetStdout(stdout).SetStderr(stderr)
 
 		t.Run("wait not yet started helper", func(t *testing.T) {
 			defer func() {
@@ -65,10 +66,8 @@ func testHelper(t *testing.T, createHelper func() helper.Helper) {
 			panic(fmt.Sprintf("unreachable: %v", h.Wait()))
 		})
 
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-
 		t.Log("starting helper stub")
-		if err := h.Start(ctx, true); err != nil {
+		if err := h.Start(true); err != nil {
 			t.Errorf("Start: error = %v", err)
 			cancel()
 			return
@@ -78,7 +77,7 @@ func testHelper(t *testing.T, createHelper func() helper.Helper) {
 
 		t.Run("start already started helper", func(t *testing.T) {
 			wantErr := "exec: already started"
-			if err := h.Start(ctx, true); err != nil && err.Error() != wantErr {
+			if err := h.Start(true); err != nil && err.Error() != wantErr {
 				t.Errorf("Start: error = %v, wantErr %v",
 					err, wantErr)
 				return
@@ -107,14 +106,14 @@ func testHelper(t *testing.T, createHelper func() helper.Helper) {
 	})
 
 	t.Run("start helper and wait", func(t *testing.T) {
-		h := createHelper()
-
-		stdout, stderr := new(strings.Builder), new(strings.Builder)
-		h.Stdout(stdout).Stderr(stderr)
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
+		h := createHelper(ctx)
+
+		stdout, stderr := new(strings.Builder), new(strings.Builder)
+		h.SetStdout(stdout).SetStderr(stderr)
 
-		if err := h.Start(ctx, false); err != nil {
+		if err := h.Start(false); err != nil {
 			t.Errorf("Start() error = %v",
 				err)
 			return
diff --git a/internal/app/shim/main.go b/internal/app/shim/main.go
index 360c89b7..4ec3916c 100644
--- a/internal/app/shim/main.go
+++ b/internal/app/shim/main.go
@@ -124,7 +124,11 @@ func Main() {
 	if fmsg.Load() {
 		seccomp.CPrintln = log.Println
 	}
+
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop() // unreachable
 	if b, err := helper.NewBwrap(
+		ctx,
 		conf, path.Join(fst.Tmp, "sbin/init0"), false,
 		nil, func(int, int) []string { return make([]string, 0) },
 		extraFiles,
@@ -132,12 +136,10 @@ func Main() {
 	); err != nil {
 		log.Fatalf("malformed sandbox config: %v", err)
 	} else {
-		b.Stdin(os.Stdin).Stdout(os.Stdout).Stderr(os.Stderr)
-		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-		defer stop() // unreachable
+		b.SetStdin(os.Stdin).SetStdout(os.Stdout).SetStderr(os.Stderr)
 
 		// run and pass through exit code
-		if err = b.Start(ctx, false); err != nil {
+		if err = b.Start(false); err != nil {
 			log.Fatalf("cannot start target process: %v", err)
 		} else if err = b.Wait(); err != nil {
 			var exitError *exec.ExitError