fix: add preExec to Run (#368)

Fixes #364

Adds the capability for Run() to run defers defined outside of Run().
Currently the only use-case for this is to close the Coder logger.
Also adds an integration test that validates this behaviour and ensures that closeLogs gets called.

Co-authored-by: Mathias Fredriksson <mafredri@gmail.com>

Author: johnstcn

cmd/envbuilder/main.go

@@ -37,6 +37,12 @@ func envbuilderCmd() serpent.Command {
 		Options: o.CLI(),
 		Handler: func(inv *serpent.Invocation) error {
 			o.SetDefaults()
+			var preExec []func()
+			defer func() { // Ensure cleanup in case of error.
+				for _, fn := range preExec {
+					fn()
+				}
+			}()
 			o.Logger = log.New(os.Stderr, o.Verbose)
 			if o.CoderAgentURL != "" {
 				if o.CoderAgentToken == "" {
@@ -50,6 +56,10 @@ func envbuilderCmd() serpent.Command {
 				if err == nil {
 					o.Logger = log.Wrap(o.Logger, coderLog)
 					defer closeLogs()
+					preExec = append(preExec, func() {
+						o.Logger(log.LevelInfo, "Closing logs")
+						closeLogs()
+					})
 					// This adds the envbuilder subsystem.
 					// If telemetry is enabled in a Coder deployment,
 					// this will be reported and help us understand
@@ -78,7 +88,7 @@ func envbuilderCmd() serpent.Command {
 				return nil
 			}
 
-			err := envbuilder.Run(inv.Context(), o)
+			err := envbuilder.Run(inv.Context(), o, preExec...)
 			if err != nil {
 				o.Logger(log.LevelError, "error: %s", err)
 			}

envbuilder.go

@@ -84,7 +84,9 @@ type execArgsInfo struct {
 // Logger is the logf to use for all operations.
 // Filesystem is the filesystem to use for all operations.
 // Defaults to the host filesystem.
-func Run(ctx context.Context, opts options.Options) error {
+// preExec are any functions that should be called before exec'ing the init
+// command. This is useful for ensuring that defers get run.
+func Run(ctx context.Context, opts options.Options, preExec ...func()) error {
 	var args execArgsInfo
 	// Run in a separate function to ensure all defers run before we
 	// setuid or exec.
@@ -103,6 +105,9 @@ func Run(ctx context.Context, opts options.Options) error {
 	}
 
 	opts.Logger(log.LevelInfo, "=== Running the init command %s %+v as the %q user...", opts.InitCommand, args.InitArgs, args.UserInfo.user.Username)
+	for _, fn := range preExec {
+		fn()
+	}
 
 	err = syscall.Exec(args.InitCommand, append([]string{args.InitCommand}, args.InitArgs...), args.Environ)
 	if err != nil {

integration/integration_test.go

@@ -23,6 +23,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/envbuilder"
 	"github.com/coder/envbuilder/devcontainer/features"
 	"github.com/coder/envbuilder/internal/magicdir"
@@ -58,6 +60,71 @@ const (
 	testImageUbuntu    = "localhost:5000/envbuilder-test-ubuntu:latest"
 )
 
+func TestLogs(t *testing.T) {
+	t.Parallel()
+
+	token := uuid.NewString()
+	logsDone := make(chan struct{})
+
+	logHandler := func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/api/v2/buildinfo":
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"version": "v2.8.9"}`))
+			return
+		case "/api/v2/workspaceagents/me/logs":
+			w.WriteHeader(http.StatusOK)
+			tokHdr := r.Header.Get(codersdk.SessionTokenHeader)
+			assert.Equal(t, token, tokHdr)
+			var req agentsdk.PatchLogs
+			err := json.NewDecoder(r.Body).Decode(&req)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return
+			}
+			for _, log := range req.Logs {
+				t.Logf("got log: %+v", log)
+				if strings.Contains(log.Output, "Closing logs") {
+					close(logsDone)
+					return
+				}
+			}
+			return
+		default:
+			t.Errorf("unexpected request to %s", r.URL.Path)
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+	}
+	logSrv := httptest.NewServer(http.HandlerFunc(logHandler))
+	defer logSrv.Close()
+
+	// Ensures that a Git repository with a devcontainer.json is cloned and built.
+	srv := gittest.CreateGitServer(t, gittest.Options{
+		Files: map[string]string{
+			"devcontainer.json": `{
+				"build": {
+					"dockerfile": "Dockerfile"
+				},
+			}`,
+			"Dockerfile": fmt.Sprintf(`FROM %s`, testImageUbuntu),
+		},
+	})
+	_, err := runEnvbuilder(t, runOpts{env: []string{
+		envbuilderEnv("GIT_URL", srv.URL),
+		"CODER_AGENT_URL=" + logSrv.URL,
+		"CODER_AGENT_TOKEN=" + token,
+	}})
+	require.NoError(t, err)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timed out waiting for logs")
+	case <-logsDone:
+	}
+}
+
 func TestInitScriptInitCommand(t *testing.T) {
 	t.Parallel()