exec: allow specifying non-zero exit codes for execs

client/client_test.go

@@ -220,6 +220,7 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testOCIIndexMediatype,
 	testLayerLimitOnMounts,
 	testFrontendVerifyPlatforms,
+	testRunValidExitCodes,
 }
 
 func TestIntegration(t *testing.T) {
@@ -10588,3 +10589,47 @@ func testClientCustomGRPCOpts(t *testing.T, sb integration.Sandbox) {
 
 	require.Contains(t, interceptedMethods, "/moby.buildkit.v1.Control/Solve")
 }
+
+func testRunValidExitCodes(t *testing.T, sb integration.Sandbox) {
+	requiresLinux(t)
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	// no exit codes specified, equivalent to [0]
+	out := llb.Image("busybox:latest")
+	out = out.Run(llb.Shlex(`sh -c "exit 0"`)).Root()
+	out = out.Run(llb.Shlex(`sh -c "exit 1"`)).Root()
+	def, err := out.Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.Error(t, err)
+	require.ErrorContains(t, err, "exit code: 1")
+
+	// empty exit codes, equivalent to [0]
+	out = llb.Image("busybox:latest")
+	out = out.Run(llb.Shlex(`sh -c "exit 0"`), llb.ValidExitCodes()).Root()
+	def, err = out.Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	// if we expect non-zero, those non-zero codes should succeed
+	out = llb.Image("busybox:latest")
+	out = out.Run(llb.Shlex(`sh -c "exit 1"`), llb.ValidExitCodes(1)).Root()
+	out = out.Run(llb.Shlex(`sh -c "exit 2"`), llb.ValidExitCodes(2, 3)).Root()
+	out = out.Run(llb.Shlex(`sh -c "exit 3"`), llb.ValidExitCodes(2, 3)).Root()
+	def, err = out.Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	// if we expect non-zero, returning zero should fail
+	out = llb.Image("busybox:latest")
+	out = out.Run(llb.Shlex(`sh -c "exit 0"`), llb.ValidExitCodes(1)).Root()
+	def, err = out.Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.Error(t, err)
+	require.ErrorContains(t, err, "exit code: 0")
+}

client/llb/exec.go

@@ -193,6 +193,17 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		return "", nil, nil, nil, err
 	}
 
+	var validExitCodes []int32
+	if codes, err := getValidExitCodes(e.base)(ctx, c); err != nil {
+		return "", nil, nil, nil, err
+	} else if codes != nil {
+		validExitCodes = make([]int32, len(codes))
+		for i, code := range codes {
+			validExitCodes[i] = int32(code)
+		}
+		addCap(&e.constraints, pb.CapExecValidExitCode)
+	}
+
 	meta := &pb.Meta{
 		Args:                      args,
 		Env:                       env.ToArray(),
@@ -201,6 +212,7 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		Hostname:                  hostname,
 		CgroupParent:              cgrpParent,
 		RemoveMountStubsRecursive: true,
+		ValidExitCodes:            validExitCodes,
 	}
 
 	extraHosts, err := getExtraHosts(e.base)(ctx, c)
@@ -581,6 +593,7 @@ func Shlex(str string) RunOption {
 		ei.State = shlexf(str, false)(ei.State)
 	})
 }
+
 func Shlexf(str string, v ...interface{}) RunOption {
 	return runOptionFunc(func(ei *ExecInfo) {
 		ei.State = shlexf(str, true, v...)(ei.State)
@@ -605,6 +618,12 @@ func AddUlimit(name UlimitName, soft int64, hard int64) RunOption {
 	})
 }
 
+func ValidExitCodes(codes ...int) RunOption {
+	return runOptionFunc(func(ei *ExecInfo) {
+		ei.State = validExitCodes(codes...)(ei.State)
+	})
+}
+
 func WithCgroupParent(cp string) RunOption {
 	return runOptionFunc(func(ei *ExecInfo) {
 		ei.State = ei.State.WithCgroupParent(cp)

client/llb/meta.go

@@ -18,14 +18,15 @@ import (
 type contextKeyT string
 
 var (
-	keyArgs         = contextKeyT("llb.exec.args")
-	keyDir          = contextKeyT("llb.exec.dir")
-	keyEnv          = contextKeyT("llb.exec.env")
-	keyExtraHost    = contextKeyT("llb.exec.extrahost")
-	keyHostname     = contextKeyT("llb.exec.hostname")
-	keyUlimit       = contextKeyT("llb.exec.ulimit")
-	keyCgroupParent = contextKeyT("llb.exec.cgroup.parent")
-	keyUser         = contextKeyT("llb.exec.user")
+	keyArgs           = contextKeyT("llb.exec.args")
+	keyDir            = contextKeyT("llb.exec.dir")
+	keyEnv            = contextKeyT("llb.exec.env")
+	keyExtraHost      = contextKeyT("llb.exec.extrahost")
+	keyHostname       = contextKeyT("llb.exec.hostname")
+	keyUlimit         = contextKeyT("llb.exec.ulimit")
+	keyCgroupParent   = contextKeyT("llb.exec.cgroup.parent")
+	keyUser           = contextKeyT("llb.exec.user")
+	keyValidExitCodes = contextKeyT("llb.exec.validexitcodes")
 
 	keyPlatform = contextKeyT("llb.platform")
 	keyNetwork  = contextKeyT("llb.network")
@@ -165,6 +166,25 @@ func getUser(s State) func(context.Context, *Constraints) (string, error) {
 	}
 }
 
+func validExitCodes(codes ...int) StateOption {
+	return func(s State) State {
+		return s.WithValue(keyValidExitCodes, codes)
+	}
+}
+
+func getValidExitCodes(s State) func(context.Context, *Constraints) ([]int, error) {
+	return func(ctx context.Context, c *Constraints) ([]int, error) {
+		v, err := s.getValue(keyValidExitCodes)(ctx, c)
+		if err != nil {
+			return nil, err
+		}
+		if v != nil {
+			return v.([]int), nil
+		}
+		return nil, nil
+	}
+}
+
 // Hostname returns a [StateOption] which sets the hostname used for containers created by [State.Run].
 // This is the equivalent of [State.Hostname]
 // See [State.With] for where to use this.
@@ -312,6 +332,7 @@ func Network(v pb.NetMode) StateOption {
 		return s.WithValue(keyNetwork, v)
 	}
 }
+
 func getNetwork(s State) func(context.Context, *Constraints) (pb.NetMode, error) {
 	return func(ctx context.Context, c *Constraints) (pb.NetMode, error) {
 		v, err := s.getValue(keyNetwork)(ctx, c)
@@ -334,6 +355,7 @@ func Security(v pb.SecurityMode) StateOption {
 		return s.WithValue(keySecurity, v)
 	}
 }
+
 func getSecurity(s State) func(context.Context, *Constraints) (pb.SecurityMode, error) {
 	return func(ctx context.Context, c *Constraints) (pb.SecurityMode, error) {
 		v, err := s.getValue(keySecurity)(ctx, c)

executor/containerdexecutor/executor.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"slices"
 	"sync"
 	"syscall"
 	"time"
@@ -215,7 +216,7 @@ func (w *containerdExecutor) Run(ctx context.Context, id string, root executor.M
 	}
 
 	trace.SpanFromContext(ctx).AddEvent("Container created")
-	err = w.runProcess(ctx, task, process.Resize, process.Signal, func() {
+	err = w.runProcess(ctx, task, process.Resize, process.Signal, process.Meta.ValidExitCodes, func() {
 		startedOnce.Do(func() {
 			trace.SpanFromContext(ctx).AddEvent("Container started")
 			if started != nil {
@@ -306,7 +307,7 @@ func (w *containerdExecutor) Exec(ctx context.Context, id string, process execut
 		return errors.WithStack(err)
 	}
 
-	err = w.runProcess(ctx, taskProcess, process.Resize, process.Signal, nil)
+	err = w.runProcess(ctx, taskProcess, process.Resize, process.Signal, process.Meta.ValidExitCodes, nil)
 	return err
 }
 
@@ -323,7 +324,7 @@ func fixProcessOutput(process *executor.ProcessInfo) {
 	}
 }
 
-func (w *containerdExecutor) runProcess(ctx context.Context, p containerd.Process, resize <-chan executor.WinSize, signal <-chan syscall.Signal, started func()) error {
+func (w *containerdExecutor) runProcess(ctx context.Context, p containerd.Process, resize <-chan executor.WinSize, signal <-chan syscall.Signal, validExitCodes []int, started func()) error {
 	// Not using `ctx` here because the context passed only affects the statusCh which we
 	// don't want cancelled when ctx.Done is sent.  We want to process statusCh on cancel.
 	statusCh, err := p.Wait(context.Background())
@@ -408,22 +409,30 @@ func (w *containerdExecutor) runProcess(ctx context.Context, p containerd.Proces
 					attribute.Int("exit.code", int(status.ExitCode())),
 				),
 			)
-			if status.ExitCode() != 0 {
-				exitErr := &gatewayapi.ExitError{
-					ExitCode: status.ExitCode(),
-					Err:      status.Error(),
-				}
-				if status.ExitCode() == gatewayapi.UnknownExitStatus && status.Error() != nil {
-					exitErr.Err = errors.Wrap(status.Error(), "failure waiting for process")
-				}
-				select {
-				case <-ctx.Done():
-					exitErr.Err = errors.Wrap(context.Cause(ctx), exitErr.Error())
-				default:
+
+			if validExitCodes == nil {
+				// no exit codes specified, so only 0 is allowed
+				if status.ExitCode() == 0 {
+					return nil
 				}
-				return exitErr
+			} else if slices.Contains(validExitCodes, int(status.ExitCode())) {
+				// exit code in allowed list, so exit cleanly
+				return nil
+			}
+
+			exitErr := &gatewayapi.ExitError{
+				ExitCode: status.ExitCode(),
+				Err:      status.Error(),
+			}
+			if status.ExitCode() == gatewayapi.UnknownExitStatus && status.Error() != nil {
+				exitErr.Err = errors.Wrap(status.Error(), "failure waiting for process")
+			}
+			select {
+			case <-ctx.Done():
+				exitErr.Err = errors.Wrap(context.Cause(ctx), exitErr.Error())
+			default:
 			}
-			return nil
+			return exitErr
 		case <-killCtxDone:
 			if cancel != nil {
 				cancel(errors.WithStack(context.Canceled))

executor/executor.go

@@ -25,6 +25,7 @@ type Meta struct {
 	CgroupParent   string
 	NetMode        pb.NetMode
 	SecurityMode   pb.SecurityMode
+	ValidExitCodes []int
 
 	RemoveMountStubsRecursive bool
 }

executor/runcexecutor/executor.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"sync"
 	"syscall"
@@ -335,7 +336,7 @@ func (w *runcExecutor) Run(ctx context.Context, id string, root executor.Mount,
 	}
 	doReleaseNetwork = false
 
-	err = exitError(ctx, cgroupPath, err)
+	err = exitError(ctx, cgroupPath, err, process.Meta.ValidExitCodes)
 	if err != nil {
 		if rec != nil {
 			rec.Close()
@@ -351,41 +352,44 @@ func (w *runcExecutor) Run(ctx context.Context, id string, root executor.Mount,
 	return rec, rec.CloseAsync(releaseContainer)
 }
 
-func exitError(ctx context.Context, cgroupPath string, err error) error {
-	if err != nil {
-		exitErr := &gatewayapi.ExitError{
-			ExitCode: gatewayapi.UnknownExitStatus,
-			Err:      err,
-		}
+func exitError(ctx context.Context, cgroupPath string, err error, validExitCodes []int) error {
+	exitErr := &gatewayapi.ExitError{ExitCode: uint32(gatewayapi.UnknownExitStatus), Err: err}
+
+	if err == nil {
+		exitErr.ExitCode = 0
+	} else {
 		var runcExitError *runc.ExitError
-		if errors.As(err, &runcExitError) && runcExitError.Status >= 0 {
-			exitErr = &gatewayapi.ExitError{
-				ExitCode: uint32(runcExitError.Status),
-			}
+		if errors.As(err, &runcExitError) {
+			exitErr = &gatewayapi.ExitError{ExitCode: uint32(runcExitError.Status)}
 		}
 
 		detectOOM(ctx, cgroupPath, exitErr)
-
-		trace.SpanFromContext(ctx).AddEvent(
-			"Container exited",
-			trace.WithAttributes(
-				attribute.Int("exit.code", int(exitErr.ExitCode)),
-			),
-		)
-		select {
-		case <-ctx.Done():
-			exitErr.Err = errors.Wrap(context.Cause(ctx), exitErr.Error())
-			return exitErr
-		default:
-			return stack.Enable(exitErr)
-		}
 	}
 
 	trace.SpanFromContext(ctx).AddEvent(
 		"Container exited",
-		trace.WithAttributes(attribute.Int("exit.code", 0)),
+		trace.WithAttributes(attribute.Int("exit.code", int(exitErr.ExitCode))),
 	)
-	return nil
+
+	if validExitCodes == nil {
+		// no exit codes specified, so only 0 is allowed
+		if exitErr.ExitCode == 0 {
+			return nil
+		}
+	} else {
+		// exit code in allowed list, so exit cleanly
+		if slices.Contains(validExitCodes, int(exitErr.ExitCode)) {
+			return nil
+		}
+	}
+
+	select {
+	case <-ctx.Done():
+		exitErr.Err = errors.Wrap(context.Cause(ctx), exitErr.Error())
+		return exitErr
+	default:
+		return stack.Enable(exitErr)
+	}
 }
 
 func (w *runcExecutor) Exec(ctx context.Context, id string, process executor.ProcessInfo) (err error) {
@@ -456,7 +460,7 @@ func (w *runcExecutor) Exec(ctx context.Context, id string, process executor.Pro
 	}
 
 	err = w.exec(ctx, id, spec.Process, process, nil)
-	return exitError(ctx, "", err)
+	return exitError(ctx, "", err, process.Meta.ValidExitCodes)
 }
 
 type forwardIO struct {

solver/llbsolver/ops/exec.go

@@ -472,6 +472,13 @@ func (e *ExecOp) Exec(ctx context.Context, g session.Group, inputs []solver.Resu
 	}
 	meta.Env = append(meta.Env, secretEnv...)
 
+	if e.op.Meta.ValidExitCodes != nil {
+		meta.ValidExitCodes = make([]int, len(e.op.Meta.ValidExitCodes))
+		for i, code := range e.op.Meta.ValidExitCodes {
+			meta.ValidExitCodes[i] = int(code)
+		}
+	}
+
 	stdout, stderr, flush := logs.NewLogStreams(ctx, os.Getenv("BUILDKIT_DEBUG_EXEC_OUTPUT") == "1")
 	defer stdout.Close()
 	defer stderr.Close()

solver/pb/caps.go

@@ -60,6 +60,7 @@ const (
 	CapExecMountContentCache             apicaps.CapID = "exec.mount.cache.content"
 	CapExecCgroupsMounted                apicaps.CapID = "exec.cgroup"
 	CapExecSecretEnv                     apicaps.CapID = "exec.secretenv"
+	CapExecValidExitCode                 apicaps.CapID = "exec.validexitcode"
 
 	CapFileBase                               apicaps.CapID = "file.base"
 	CapFileRmWildcard                         apicaps.CapID = "file.rm.wildcard"
@@ -357,6 +358,12 @@ func init() {
 		Status:  apicaps.CapStatusExperimental,
 	})
 
+	Caps.Init(apicaps.Cap{
+		ID:      CapExecValidExitCode,
+		Enabled: true,
+		Status:  apicaps.CapStatusExperimental,
+	})
+
 	Caps.Init(apicaps.Cap{
 		ID:      CapFileBase,
 		Enabled: true,