[ONPREM-2408] Add entrypoint override (#174)

* Add entrypoint override

* Add changelog entry

* Make entrypoint override a command

* Fix entrypoint args

* Add acceptance test

* Skip acceptance test on Windows

* Add better help description for entrypoint override

Author: christian-stephen

CHANGELOG.md

@@ -10,6 +10,7 @@ By following these guidelines, we can easily determine which changes should be i
 
 ## Edge
 
+- [#174](https://github.com/circleci/runner-init/pull/174) Add support for a custom entrypoint override.
 - [#224](https://github.com/circleci/runner-init/pull/224) Record and log timings in orchestrator init function.
 - [#212](https://github.com/circleci/runner-init/pull/212) Fix child process cleanup on Windows using job objects. This ensures that child processes are destroyed when the parent process (task-agent) terminates.
 - [#197](https://github.com/circleci/runner-init/pull/197) Fix `%PATH%` on Windows by using the OS-specific path list separator.

acceptance/acceptance_test.go

@@ -12,8 +12,9 @@ import (
 )
 
 var (
-	orchestratorTestBinary        = os.Getenv("ORCHESTRATOR_TEST_BINARY")
-	orchestratorTestBinaryRunTask = ""
+	orchestratorTestBinary         = os.Getenv("ORCHESTRATOR_TEST_BINARY")
+	orchestratorTestBinaryRunTask  = ""
+	orchestratorTestBinaryOverride = ""
 
 	binariesPath    = ""
 	taskAgentBinary = ""
@@ -60,32 +61,35 @@ func runTests(m *testing.M) (int, error) {
 	fmt.Printf("Using 'orchestrator' test binary: %q\n", orchestratorTestBinary)
 	fmt.Printf("Using fake 'task-agent' test binary: %q\n", taskAgentBinary)
 
-	if err := createRunTaskScript(); err != nil {
+	if orchestratorTestBinaryRunTask, err = createRunnerScript("run-task"); err != nil {
 		return 0, err
 	}
 	fmt.Printf("Using 'orchestrator run-task' script: %q\n", orchestratorTestBinaryRunTask)
 
+	if orchestratorTestBinaryOverride, err = createRunnerScript("override"); err != nil {
+		return 0, err
+	}
+	fmt.Printf("Using 'orchestrator override' script: %q\n", orchestratorTestBinaryOverride)
+
 	return m.Run(), nil
 }
 
 // A little hack to get around limitations of the test runner on positional arguments
-func createRunTaskScript() error {
+func createRunnerScript(cmd string) (string, error) {
 	var script string
 	var scriptPath string
 
 	if runtime.GOOS == "windows" {
-		script = "@echo off\n" + orchestratorTestBinary + " run-task"
-		scriptPath = filepath.Join(binariesPath, "orchestratorRunTask.bat")
+		script = "@echo off\n" + orchestratorTestBinary + " " + cmd
+		scriptPath = filepath.Join(binariesPath, fmt.Sprintf("orchestrator-%s.bat", cmd))
 	} else {
-		script = "#!/bin/bash\nexec " + orchestratorTestBinary + " run-task"
-		scriptPath = filepath.Join(binariesPath, "orchestratorRunTask.sh")
+		script = "#!/bin/bash\nexec " + orchestratorTestBinary + " " + cmd
+		scriptPath = filepath.Join(binariesPath, fmt.Sprintf("orchestrator-%s.sh", cmd))
 	}
 
 	if err := os.WriteFile(scriptPath, []byte(script), 0750); err != nil { //nolint:gosec
-		return err
+		return "", err
 	}
 
-	orchestratorTestBinaryRunTask = scriptPath
-
-	return nil
+	return scriptPath, nil
 }

acceptance/task_test.go

@@ -4,12 +4,14 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/circleci/ex/testing/runner"
 	"gotest.tools/v3/assert"
+	"gotest.tools/v3/assert/cmp"
 )
 
 func TestRunTask(t *testing.T) {
@@ -26,32 +28,81 @@ func TestRunTask(t *testing.T) {
 	"max_run_time": 60000000000
 }`, strings.ReplaceAll(readinessFilePath, `\`, `\\`), strings.ReplaceAll(taskAgentBinary, `\`, `\\`))
 
-	r := runner.New(
-		"CIRCLECI_GOAT_SHUTDOWN_DELAY=10s",
-		"CIRCLECI_GOAT_CONFIG="+goodConfig,
-		"CIRCLECI_GOAT_HEALTH_CHECK_ADDR=:7624",
-	)
-	res, err := r.Start(orchestratorTestBinaryRunTask)
-	assert.NilError(t, err)
+	t.Run("Good run-task", func(t *testing.T) {
+		r := runner.New(
+			"CIRCLECI_GOAT_SHUTDOWN_DELAY=10s",
+			"CIRCLECI_GOAT_CONFIG="+goodConfig,
+			"CIRCLECI_GOAT_HEALTH_CHECK_ADDR=:7624",
+		)
+		res, err := r.Start(orchestratorTestBinaryRunTask)
+		assert.NilError(t, err)
+
+		t.Run("Probe for readiness", func(t *testing.T) {
+			assert.NilError(t, res.Ready("admin", time.Second*20))
+		})
+
+		go func() {
+			f, err := os.Create(readinessFilePath) //nolint:gosec
+			defer func() { assert.NilError(t, f.Close()) }()
+			assert.NilError(t, err)
+		}()
 
-	t.Run("Probe for readiness", func(t *testing.T) {
-		assert.NilError(t, res.Ready("admin", time.Second*20))
+		t.Run("Run task", func(t *testing.T) {
+			select {
+			case err = <-res.Wait():
+				assert.NilError(t, err)
+			case <-time.After(time.Second * 40):
+				assert.NilError(t, res.Stop())
+				t.Fatal(t, "timeout before process stopped")
+			}
+		})
 	})
 
-	go func() {
-		f, err := os.Create(readinessFilePath) //nolint:gosec
-		defer func() { assert.NilError(t, f.Close()) }()
+	t.Run("Good entrypoint override", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("Not supported on Windows")
+		}
+
+		entrypointPath := filepath.ToSlash(filepath.Join(t.TempDir(), "entrypoint.sh"))
+
+		//nolint:gosec
+		err := os.WriteFile(entrypointPath, []byte(`#!/bin/bash
+echo "Executing custom entrypoint"
+exec "$@"`), 0750)
 		assert.NilError(t, err)
-	}()
 
-	t.Run("Run task", func(t *testing.T) {
-		select {
-		case err = <-res.Wait():
+		r := runner.New(
+			"CIRCLECI_GOAT_ENTRYPOINT="+entrypointPath,
+			"CIRCLECI_GOAT_SHUTDOWN_DELAY=10s",
+			"CIRCLECI_GOAT_CONFIG="+goodConfig,
+			"CIRCLECI_GOAT_HEALTH_CHECK_ADDR=:7624",
+		)
+		res, err := r.Start(orchestratorTestBinaryOverride)
+		assert.NilError(t, err)
+
+		t.Run("Probe for readiness", func(t *testing.T) {
+			assert.NilError(t, res.Ready("admin", time.Second*20))
+		})
+
+		t.Run("Custom entrypoint ran", func(t *testing.T) {
+			assert.Check(t, cmp.Contains(res.Logs(), "Executing custom entrypoint"))
+		})
+
+		go func() {
+			f, err := os.Create(readinessFilePath) //nolint:gosec
+			defer func() { assert.NilError(t, f.Close()) }()
 			assert.NilError(t, err)
-		case <-time.After(time.Second * 40):
-			assert.NilError(t, res.Stop())
-			t.Fatal(t, "timeout before process stopped")
-		}
+		}()
+
+		t.Run("Run task", func(t *testing.T) {
+			select {
+			case err = <-res.Wait():
+				assert.NilError(t, err)
+			case <-time.After(time.Second * 40):
+				assert.NilError(t, res.Stop())
+				t.Fatal(t, "timeout before process stopped")
+			}
+		})
 	})
 
 	// TODO: Add more test cases...

cmd/orchestrator/help_test.go

@@ -34,6 +34,11 @@ func TestHelp(t *testing.T) {
 			cli:          &cli.Init,
 			wantFilename: "init.txt",
 		},
+		{
+			name:         "check override command help",
+			cli:          &cli.Override,
+			wantFilename: "override.txt",
+		},
 		{
 			name:         "check run-task command help",
 			cli:          &cli.RunTask,

cmd/orchestrator/main.go

@@ -19,14 +19,16 @@ import (
 	"github.com/circleci/runner-init/cmd/setup"
 	initialize "github.com/circleci/runner-init/init"
 	"github.com/circleci/runner-init/task"
+	"github.com/circleci/runner-init/task/entrypoint"
 	"github.com/circleci/runner-init/task/taskerrors"
 )
 
 type cli struct {
 	Version kong.VersionFlag `short:"v" help:"Print version information and quit."`
 
-	Init    initCmd    `cmd:"" name:"init" default:"withargs"`
-	RunTask runTaskCmd `cmd:"" name:"run-task"`
+	Init     initCmd     `cmd:"" name:"init" default:"withargs"`
+	Override overrideCmd `cmd:"" name:"override"`
+	RunTask  runTaskCmd  `cmd:"" name:"run-task"`
 
 	ShutdownDelay time.Duration `default:"0s" help:"Delay shutdown by this amount."`
 }
@@ -36,6 +38,12 @@ type initCmd struct {
 	Destination string `arg:"" env:"DESTINATION" type:"path" default:"/opt/circleci/bin" help:"Path where to copy the agent binaries to."`
 }
 
+type overrideCmd struct {
+	Entrypoint []string `help:"Custom init process to execute as PID 1, overriding orchestrator. Must accept and execute the orchestrator command/arguments (e.g., exec \"$@\"), propagate signals, and handle standard init responsibilities like reaping zombie processes."`
+
+	runTaskCmd
+}
+
 type runTaskCmd struct {
 	TerminationGracePeriod time.Duration `default:"10s" help:"How long the agent will wait for the task to complete if interrupted."`
 	HealthCheckAddr        string        `default:":7623" help:"Address for the health check API to listen on."`
@@ -87,6 +95,13 @@ func run(version, date string) (err error) {
 			return initialize.Run(ctx, c.Source, c.Destination)
 		})
 
+	case "override":
+		ep := entrypoint.New(cli.Override.Entrypoint)
+		sys.AddService(func(ctx context.Context) error {
+			defer cancel()
+			return ep.Run(ctx)
+		})
+
 	case "run-task":
 		orchestrator, err := runSetup(ctx, cli, version, sys)
 		if err != nil {
@@ -102,9 +117,8 @@ func run(version, date string) (err error) {
 	return sys.Run(ctx, cli.ShutdownDelay)
 }
 
-func runSetup(ctx context.Context, cli cli, version string, sys *system.System) (*task.Orchestrator, error) {
+func runSetup(ctx context.Context, cli cli, version string, sys *system.System) (Runner, error) {
 	c := cli.RunTask
-
 	// Strip the orchestrator configuration from the environment
 	_ = os.Unsetenv("CIRCLECI_GOAT_CONFIG")
 
@@ -130,3 +144,7 @@ func runSetup(ctx context.Context, cli cli, version string, sys *system.System)
 
 	return o, nil
 }
+
+type Runner interface {
+	Run(ctx context.Context) error
+}

cmd/orchestrator/testdata/help.txt

@@ -10,6 +10,8 @@ Flags:
 Commands:
   init [<source> [<destination>]] [flags]
 
+  override [flags]
+
   run-task [flags]
 
 Run "test-app <command> --help" for more information on a command.

cmd/orchestrator/testdata/override.txt

@@ -0,0 +1,16 @@
+Usage: test-app [flags]
+
+Flags:
+  -h, --help    Show context-sensitive help.
+      --entrypoint=ENTRYPOINT,...
+                Custom init process to execute as PID 1, overriding
+                orchestrator. Must accept and execute the orchestrator
+                command/arguments (e.g., exec "$@"), propagate signals,
+                and handle standard init responsibilities like reaping zombie
+                processes ($CIRCLECI_GOAT_ENTRYPOINT).
+      --termination-grace-period=10s
+                How long the agent will wait for the task to complete if
+                interrupted ($CIRCLECI_GOAT_TERMINATION_GRACE_PERIOD).
+      --health-check-addr=":7623"
+                Address for the health check API to listen on
+                ($CIRCLECI_GOAT_HEALTH_CHECK_ADDR).

task/cmd/cmd.go

@@ -6,7 +6,6 @@ import (
 	"io"
 	"os"
 	"os/exec"
-	"strings"
 	"sync/atomic"
 	"syscall"
 )
@@ -106,17 +105,7 @@ func newCmd(ctx context.Context, argv []string, user string, stderrSaver *prefix
 	//#nosec:G204 // this is intentionally setting up a command
 	cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
 
-	for _, env := range os.Environ() {
-		if strings.HasPrefix(env, "CIRCLECI_GOAT") {
-			// Prevent internal configuration from being injected in the command environment
-			continue
-		}
-		cmd.Env = append(cmd.Env, env)
-	}
-	if env != nil {
-		cmd.Env = append(cmd.Env, env...)
-	}
-
+	cmd.Env = Environ(env...)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = io.MultiWriter(os.Stderr, stderrSaver)
 

task/cmd/environ.go

@@ -0,0 +1,21 @@
+package cmd
+
+import (
+	"os"
+	"strings"
+)
+
+func Environ(extraEnv ...string) (environ []string) {
+	for _, env := range os.Environ() {
+		if strings.HasPrefix(env, "CIRCLECI_GOAT") {
+			// Prevent internal configuration from being unintentionally injected in the command environment
+			continue
+		}
+		environ = append(environ, env)
+	}
+	if extraEnv != nil {
+		environ = append(environ, extraEnv...)
+	}
+
+	return environ
+}

task/entrypoint/entrypoint.go

@@ -0,0 +1,35 @@
+package entrypoint
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"syscall"
+
+	"github.com/circleci/ex/o11y"
+)
+
+type Entrypoint struct {
+	args []string
+}
+
+func New(args []string) Entrypoint {
+	return Entrypoint{args}
+}
+
+func (e Entrypoint) Run(ctx context.Context) (err error) {
+	_, span := o11y.StartSpan(ctx, "override-entrypoint")
+	defer o11y.End(span, &err)
+
+	args := append([]string{e.args[0]}, os.Args[0], "run-task")
+	if len(e.args) > 1 {
+		args = append(args, e.args[1:]...)
+	}
+
+	//#nosec:G204 // this is intentionally setting up a command
+	if err := syscall.Exec(e.args[0], args, os.Environ()); err != nil {
+		return fmt.Errorf("error executing entrypoint override: %w", err)
+	}
+
+	return nil
+}