Merge pull request #1520 from buildkite/keithduncan/fix-early-exit-ho…

…ok-wd

Fix early exit hook wd change

bootstrap/bootstrap.go

@@ -411,17 +411,32 @@ func (b *Bootstrap) executeHook(ctx context.Context, scope string, name string,
 	// Get changed environment
 	changes, err := script.Changes()
 	if err != nil {
-		return errors.Wrapf(err, "Failed to get environment")
+		// Could not compute the changes in environment or working directory
+		// for some reason...
+
+		switch err.(type) {
+		case *hook.HookExitError:
+			// ...because the hook called exit(), tsk we ignore any changes
+			// since we can't discern them but continue on with the job
+			break;
+		default:
+			// ...because something else happened, report it and stop the job
+			return errors.Wrapf(err, "Failed to get environment")
+		}
+	} else {
+		// Hook exited successfully (and not early!) We have an environment and
+		// wd change we can apply to our subsequent phases
+		b.applyEnvironmentChanges(changes, redactors)
 	}
 
-	// Finally, apply changes to the current shell and config
-	b.applyEnvironmentChanges(changes, redactors)
 	return nil
 }
 
 func (b *Bootstrap) applyEnvironmentChanges(changes hook.HookScriptChanges, redactors RedactorMux) {
-	if changes.Dir != b.shell.Getwd() {
-		_ = b.shell.Chdir(changes.Dir)
+	if afterWd, err := changes.GetAfterWd(); err == nil {
+		if afterWd != b.shell.Getwd() {
+			_ = b.shell.Chdir(afterWd)
+		}
 	}
 
 	// Do we even have any environment variables to change?

bootstrap/integration/hooks_integration_test.go

@@ -58,8 +58,77 @@ func TestEnvironmentVariablesPassBetweenHooks(t *testing.T) {
 		if err := bintest.ExpectEnv(t, c.Env, `MY_CUSTOM_ENV=1`, `LLAMAS_ROCK=absolutely`); err != nil {
 			fmt.Fprintf(c.Stderr, "%v\n", err)
 			c.Exit(1)
+		} else {
+			c.Exit(0)
+		}
+	})
+
+	tester.RunAndCheck(t, "MY_CUSTOM_ENV=1")
+}
+
+func TestHooksCanUnsetEnvironmentVariables(t *testing.T) {
+	t.Parallel()
+
+	tester, err := NewBootstrapTester()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tester.Close()
+
+	if runtime.GOOS == "windows" {
+		var preCommand = []string{
+			"@echo off",
+			"set LLAMAS_ROCK=absolutely",
+		}
+		if err := ioutil.WriteFile(filepath.Join(tester.HooksDir, "pre-command.bat"),
+			[]byte(strings.Join(preCommand, "\r\n")), 0700); err != nil {
+			t.Fatal(err)
+		}
+
+		var postCommand = []string{
+			"@echo off",
+			"set LLAMAS_ROCK=",
+		}
+		if err := ioutil.WriteFile(filepath.Join(tester.HooksDir, "post-command.bat"),
+			[]byte(strings.Join(postCommand, "\n")), 0700); err != nil {
+			t.Fatal(err)
+		}
+	} else {
+		var preCommand = []string{
+			"#!/bin/bash",
+			"export LLAMAS_ROCK=absolutely",
+		}
+		if err := ioutil.WriteFile(filepath.Join(tester.HooksDir, "pre-command"),
+			[]byte(strings.Join(preCommand, "\n")), 0700); err != nil {
+			t.Fatal(err)
+		}
+
+		var postCommand = []string{
+			"#!/bin/bash",
+			"unset LLAMAS_ROCK",
+		}
+		if err := ioutil.WriteFile(filepath.Join(tester.HooksDir, "post-command"),
+			[]byte(strings.Join(postCommand, "\n")), 0700); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	tester.ExpectGlobalHook("command").Once().AndExitWith(0).AndCallFunc(func(c *bintest.Call) {
+		if c.GetEnv("LLAMAS_ROCK") != "absolutely" {
+			fmt.Fprintf(c.Stderr, "Expected command hook to have environment variable LLAMAS_ROCK be %q, got %q\n", "absolutely", c.GetEnv("LLAMAS_ROCK"))
+			c.Exit(1)
+		} else {
+			c.Exit(0)
+		}
+	})
+
+	tester.ExpectGlobalHook("pre-exit").Once().AndExitWith(0).AndCallFunc(func(c *bintest.Call) {
+		if c.GetEnv("LLAMAS_ROCK") != "" {
+			fmt.Fprintf(c.Stderr, "Expected pre-exit hook to have environment variable LLAMAS_ROCK be empty, got %q\n", c.GetEnv("LLAMAS_ROCK"))
+			c.Exit(1)
+		} else {
+			c.Exit(0)
 		}
-		c.Exit(0)
 	})
 
 	tester.RunAndCheck(t, "MY_CUSTOM_ENV=1")
@@ -93,8 +162,46 @@ func TestDirectoryPassesBetweenHooks(t *testing.T) {
 		if c.GetEnv("MY_CUSTOM_SUBDIR") != c.Dir {
 			fmt.Fprintf(c.Stderr, "Expected current dir to be %q, got %q\n", c.GetEnv("MY_CUSTOM_SUBDIR"), c.Dir)
 			c.Exit(1)
+		} else {
+			c.Exit(0)
+		}
+	})
+
+	tester.RunAndCheck(t, "MY_CUSTOM_ENV=1")
+}
+
+func TestDirectoryPassesBetweenHooksIgnoredUnderExit(t *testing.T) {
+	t.Parallel()
+
+	tester, err := NewBootstrapTester()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tester.Close()
+
+	if runtime.GOOS == "windows" {
+		t.Skip("Not implemented for windows yet")
+	}
+
+	var script = []string{
+		"#!/bin/bash",
+		"mkdir -p ./mysubdir",
+		"export MY_CUSTOM_SUBDIR=$(cd mysubdir; pwd)",
+		"cd ./mysubdir",
+		"exit 0",
+	}
+
+	if err := ioutil.WriteFile(filepath.Join(tester.HooksDir, "pre-command"), []byte(strings.Join(script, "\n")), 0700); err != nil {
+		t.Fatal(err)
+	}
+
+	tester.ExpectGlobalHook("command").Once().AndExitWith(0).AndCallFunc(func(c *bintest.Call) {
+		if c.GetEnv("BUILDKITE_BUILD_CHECKOUT_PATH") != c.Dir {
+			fmt.Fprintf(c.Stderr, "Expected current dir to be %q, got %q\n", c.GetEnv("BUILDKITE_BUILD_CHECKOUT_PATH"), c.Dir)
+			c.Exit(1)
+		} else {
+			c.Exit(0)
 		}
-		c.Exit(0)
 	})
 
 	tester.RunAndCheck(t, "MY_CUSTOM_ENV=1")
@@ -145,9 +252,9 @@ func TestReplacingCheckoutHook(t *testing.T) {
 		fmt.Fprint(c.Stderr, out)
 		if err != nil {
 			c.Exit(1)
-			return
+		} else {
+			c.Exit(0)
 		}
-		c.Exit(0)
 	})
 
 	tester.ExpectGlobalHook("pre-checkout").Once()

bootstrap/integration/plugin_integration_test.go

@@ -59,16 +59,18 @@ func TestRunningPlugins(t *testing.T) {
 		if err := bintest.ExpectEnv(t, c.Env, `MY_CUSTOM_ENV=1`, `LLAMAS_ROCK=absolutely`); err != nil {
 			fmt.Fprintf(c.Stderr, "%v\n", err)
 			c.Exit(1)
+		} else {
+			c.Exit(0)
 		}
-		c.Exit(0)
 	})
 
 	tester.ExpectGlobalHook("command").Once().AndExitWith(0).AndCallFunc(func(c *bintest.Call) {
 		if err := bintest.ExpectEnv(t, c.Env, `MY_CUSTOM_ENV=1`, `LLAMAS_ROCK=absolutely`); err != nil {
 			fmt.Fprintf(c.Stderr, "%v\n", err)
 			c.Exit(1)
+		} else {
+			c.Exit(0)
 		}
-		c.Exit(0)
 	})
 
 	tester.RunAndCheck(t, env...)

hook/scriptwrapper.go

@@ -35,12 +35,27 @@ type ScriptWrapper struct {
 	scriptFile    *os.File
 	beforeEnvFile *os.File
 	afterEnvFile  *os.File
-	beforeWd      string
 }
 
 type HookScriptChanges struct {
 	Diff env.Diff
-	Dir string
+	afterWd string
+}
+
+func (changes *HookScriptChanges) GetAfterWd() (string, error) {
+	if changes.afterWd == "" {
+		return "", fmt.Errorf("%q was not present in the hook after environment", hookWorkingDirEnv)
+	}
+
+	return changes.afterWd, nil
+}
+
+type HookExitError struct {
+	hookPath string
+}
+
+func (e *HookExitError) Error() string {
+	return fmt.Sprintf("Hook %q early exited, could not record after environment or working directory", e.hookPath)
 }
 
 // CreateScriptWrapper creates and configures a ScriptWrapper.
@@ -97,11 +112,6 @@ func CreateScriptWrapper(hookPath string) (*ScriptWrapper, error) {
 		return nil, fmt.Errorf("Failed to find absolute path to \"%s\" (%s)", wrap.hookPath, err)
 	}
 
-	wrap.beforeWd, err = os.Getwd()
-	if err != nil {
-		return nil, err
-	}
-
 	// Create the hook runner code
 	var script string
 	if isWindows && !isBashHook && !isPwshHook {
@@ -170,39 +180,30 @@ func (wrap *ScriptWrapper) Changes() (HookScriptChanges, error) {
 	}
 
 	beforeEnv := env.FromExport(string(beforeEnvContents))
-
 	afterEnv := env.FromExport(string(afterEnvContents))
+
 	if afterEnv.Length() == 0 {
-		// If the after env is completely empty it is likely because the hook
-		// ran exit(). Instead of falling over and breaking any subsequent
-		// commands, lets fall back to the original before env since we can’t
-		// extract a meaningful diff.
-		afterEnv = beforeEnv
+		return HookScriptChanges{}, &HookExitError{hookPath: wrap.hookPath}
 	}
 
 	diff := afterEnv.Diff(beforeEnv)
 
-	wd := wrap.getAfterWd(diff)
+	// Pluck the after wd from the diff before removing the key from the diff
+	afterWd := ""
+	if afterWd == "" {
+		afterWd, _ = diff.Added[hookWorkingDirEnv]
+	}
+	if afterWd == "" {
+		if change, ok := diff.Changed[hookWorkingDirEnv]; ok {
+			afterWd = change.New
+		}
+	}
 
 	diff.Remove(hookExitStatusEnv)
 	diff.Remove(hookWorkingDirEnv)
 
 	// Bash sets this, but we don't care about it
 	diff.Remove("_")
 
-	return HookScriptChanges{Diff: diff, Dir: wd}, nil
-}
-
-func (wrap *ScriptWrapper) getAfterWd(diff env.Diff) string {
-	addedVar, ok := diff.Added[hookWorkingDirEnv]
-	if ok {
-		return addedVar
-	}
-
-	changedVar, ok := diff.Changed[hookWorkingDirEnv]
-	if ok {
-		return changedVar.New
-	}
-
-	return wrap.beforeWd
+	return HookScriptChanges{Diff: diff, afterWd: afterWd}, nil
 }

hook/scriptwrapper_test.go

@@ -120,7 +120,12 @@ func TestRunningHookDetectsChangedWorkingDirectory(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	changesDir, err := filepath.EvalSymlinks(changes.Dir)
+	afterWd, err := changes.GetAfterWd()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	changesDir, err := filepath.EvalSymlinks(afterWd)
 	if err != nil {
 		t.Fatal(err)
 	}