fix: add preExec to Run

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()