Add test ensuring that pre-bootstrap hooks reject jobs

agent/integration/job_runner_integration_test.go

@@ -1,10 +1,16 @@
 package integration
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"runtime"
 	"testing"
 
 	"github.com/buildkite/agent/v3/agent"
@@ -14,7 +20,62 @@ import (
 	"github.com/buildkite/bintest/v3"
 )
 
+func TestPreBootstrapHookRefusesJob(t *testing.T) {
+	t.Parallel()
+
+	hooksDir, err := os.MkdirTemp("", "bootstrap-hooks")
+	if err != nil {
+		t.Fatalf("making bootstrap-hooks directory: %v", err)
+	}
+
+	defer os.RemoveAll(hooksDir)
+
+	mockPB := mockPreBootstrap(t, hooksDir)
+	mockPB.Expect().Once().AndCallFunc(func(c *bintest.Call) {
+		c.Exit(1) // Fail the pre-bootstrap hook
+	})
+	defer mockPB.CheckAndClose(t)
+
+	jobID := "my-job-id"
+	j := &api.Job{
+		ID:                 jobID,
+		ChunksMaxSizeBytes: 1024,
+		Env: map[string]string{
+			"BUILDKITE_COMMAND": "echo hello world",
+		},
+	}
+
+	// create a mock agent API
+	e := createTestAgentEndpoint()
+	server := e.server("my-job-id")
+	defer server.Close()
+
+	mb := mockBootstrap(t)
+	mb.Expect().NotCalled() // The bootstrap won't be called, as the pre-bootstrap hook failed
+	defer mb.CheckAndClose(t)
+
+	runJob(t, j, server, agent.AgentConfiguration{HooksPath: hooksDir}, mb)
+
+	finishRequestBody := bytes.NewBuffer(e.calls["/jobs/my-job-id/finish"][0])
+
+	var job api.Job
+	err = json.NewDecoder(finishRequestBody).Decode(&job)
+	if err != nil {
+		t.Fatalf("decoding accept request body: %v", err)
+	}
+
+	if got, want := job.ExitStatus, "-1"; got != want {
+		t.Errorf("job.ExitStatus = %q, want %q", got, want)
+	}
+
+	if got, want := job.SignalReason, "agent_refused"; got != want {
+		t.Errorf("job.SignalReason = %q, want %q", got, want)
+	}
+}
+
 func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
+	t.Parallel()
+
 	jobToken := "actually-llamas-are-only-okay"
 
 	j := &api.Job{
@@ -26,18 +87,29 @@ func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
 		},
 	}
 
-	cfg := agent.AgentConfiguration{}
-	runJob(t, j, cfg, func(c *bintest.Call) {
+	mb := mockBootstrap(t)
+	defer mb.CheckAndClose(t)
+
+	expectBootstrapAndCall(t, mb, func(c *bintest.Call) {
 		if got, want := c.GetEnv("BUILDKITE_AGENT_ACCESS_TOKEN"), jobToken; got != want {
 			t.Errorf("c.GetEnv(BUILDKITE_AGENT_ACCESS_TOKEN) = %q, want %q", got, want)
 		}
 		c.Exit(0)
 	})
+
+	// create a mock agent API
+	e := createTestAgentEndpoint()
+	server := e.server("my-job-id")
+	defer server.Close()
+
+	runJob(t, j, server, agent.AgentConfiguration{}, mb)
 }
 
 // TODO 2023-07-17: What is this testing? How is it testing it?
 // Maybe that the job runner pulls the access token from the API client? but that's all handled in the `runJob` helper...
 func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
+	t.Parallel()
+
 	j := &api.Job{
 		ID:                 "my-job-id",
 		ChunksMaxSizeBytes: 1024,
@@ -46,16 +118,27 @@ func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
 		},
 	}
 
-	cfg := agent.AgentConfiguration{}
-	runJob(t, j, cfg, func(c *bintest.Call) {
+	mb := mockBootstrap(t)
+	defer mb.CheckAndClose(t)
+
+	expectBootstrapAndCall(t, mb, func(c *bintest.Call) {
 		if got, want := c.GetEnv("BUILDKITE_AGENT_ACCESS_TOKEN"), "llamasrock"; got != want {
 			t.Errorf("c.GetEnv(BUILDKITE_AGENT_ACCESS_TOKEN) = %q, want %q", got, want)
 		}
 		c.Exit(0)
 	})
+
+	// create a mock agent API
+	e := createTestAgentEndpoint()
+	server := e.server("my-job-id")
+	defer server.Close()
+
+	runJob(t, j, server, agent.AgentConfiguration{}, mb)
 }
 
 func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
+	t.Parallel()
+
 	j := &api.Job{
 		ID:                 "my-job-id",
 		ChunksMaxSizeBytes: 1024,
@@ -65,30 +148,38 @@ func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
 		},
 	}
 
-	cfg := agent.AgentConfiguration{CommandEval: true}
-	runJob(t, j, cfg, func(c *bintest.Call) {
+	mb := mockBootstrap(t)
+	defer mb.CheckAndClose(t)
+
+	expectBootstrapAndCall(t, mb, func(c *bintest.Call) {
 		if got, want := c.GetEnv("BUILDKITE_COMMAND_EVAL"), "true"; got != want {
 			t.Errorf("c.GetEnv(BUILDKITE_COMMAND_EVAL) = %q, want %q", got, want)
 		}
 		c.Exit(0)
 	})
-}
 
-func runJob(t *testing.T, j *api.Job, cfg agent.AgentConfiguration, bootstrap func(c *bintest.Call)) {
 	// create a mock agent API
-	server := createTestAgentEndpoint(t, "my-job-id")
+	e := createTestAgentEndpoint()
+	server := e.server("my-job-id")
 	defer server.Close()
 
-	// set up a mock bootstrap that the runner will call
-	bs, err := bintest.NewMock("buildkite-agent-bootstrap")
+	runJob(t, j, server, agent.AgentConfiguration{CommandEval: true}, mb)
+}
+
+func mockBootstrap(t *testing.T) *bintest.Mock {
+	bs, err := bintest.NewMock(fmt.Sprintf("buildkite-agent-bootstrap-%s", t.Name()))
 	if err != nil {
 		t.Fatalf("bintest.NewMock() error = %v", err)
 	}
-	defer bs.CheckAndClose(t)
 
-	// execute the callback we have inside the bootstrap mock
-	bs.Expect().Once().AndExitWith(0).AndCallFunc(bootstrap)
+	return bs
+}
+
+func expectBootstrapAndCall(t *testing.T, bs *bintest.Mock, f func(*bintest.Call)) {
+	bs.Expect().Once().AndExitWith(0).AndCallFunc(f)
+}
 
+func runJob(t *testing.T, j *api.Job, server *httptest.Server, cfg agent.AgentConfiguration, bs *bintest.Mock) {
 	l := logger.Discard
 
 	// minimal metrics, this could be cleaner
@@ -119,8 +210,21 @@ func runJob(t *testing.T, j *api.Job, cfg agent.AgentConfiguration, bootstrap fu
 	}
 }
 
-func createTestAgentEndpoint(t *testing.T, jobID string) *httptest.Server {
+type testAgentEndpoint struct {
+	calls map[string][][]byte
+}
+
+func createTestAgentEndpoint() *testAgentEndpoint {
+	return &testAgentEndpoint{
+		calls: make(map[string][][]byte, 4),
+	}
+}
+
+func (t *testAgentEndpoint) server(jobID string) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		b, _ := io.ReadAll(req.Body)
+		t.calls[req.URL.Path] = append(t.calls[req.URL.Path], b)
+
 		switch req.URL.Path {
 		case "/jobs/" + jobID:
 			rw.WriteHeader(http.StatusOK)
@@ -136,3 +240,30 @@ func createTestAgentEndpoint(t *testing.T, jobID string) *httptest.Server {
 		}
 	}))
 }
+
+func mockPreBootstrap(t *testing.T, hooksDir string) *bintest.Mock {
+	mock, err := bintest.NewMock(fmt.Sprintf("buildkite-agent-pre-bootstrap-hook-%s", t.Name()))
+	if err != nil {
+		t.Fatalf("bintest.NewMock() error = %v", err)
+	}
+
+	hookScript := filepath.Join(hooksDir, "pre-bootstrap")
+	body := ""
+
+	if runtime.GOOS == "windows" {
+		body = fmt.Sprintf("@%q", mock.Path)
+		hookScript += ".bat"
+	} else {
+		body = "#!/bin/sh\n" + mock.Path
+	}
+
+	if err := os.MkdirAll(hooksDir, 0o700); err != nil {
+		t.Fatalf("creating pre-bootstrap hook mock: os.MkdirAll() error = %v", err)
+	}
+
+	if err := os.WriteFile(hookScript, []byte(body), 0o777); err != nil {
+		t.Fatalf("creating pre-bootstrap hook mock: s.WriteFile() error = %v", err)
+	}
+
+	return mock
+}

agent/job_runner.go

@@ -411,7 +411,10 @@ func (r *JobRunner) Run(ctx context.Context) error {
 	// whether it is happy to proceed.
 	shouldRunJob := true
 
+	fmt.Println("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
+	fmt.Println(r.conf.AgentConfiguration.HooksPath)
 	if hook, _ := hook.Find(r.conf.AgentConfiguration.HooksPath, "pre-bootstrap"); hook != "" {
+		fmt.Println("found pre-bootstrap hook")
 		// Once we have a hook any failure to run it MUST be fatal to the job to guarantee a true positive result from the hook
 		ok, err := r.executePreBootstrapHook(ctx, hook)
 		if !ok {
@@ -773,6 +776,7 @@ func (r *JobRunner) executePreBootstrapHook(ctx context.Context, hook string) (b
 	}
 
 	if err := sh.RunWithoutPrompt(ctx, hook); err != nil {
+		fmt.Printf("err: %s\n", err)
 		r.logger.Error("Finished pre-bootstrap hook %q: job rejected", hook)
 		return false, err
 	}

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/gofrs/flock v0.8.1
 	github.com/google/go-cmp v0.5.9
 	github.com/google/go-querystring v1.1.0
+	github.com/kr/pretty v0.2.1
 	github.com/mattn/go-zglob v0.0.4
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/oleiade/reflections v1.0.1
@@ -70,6 +71,7 @@ require (
 	github.com/googleapis/gax-go/v2 v2.11.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.15.2 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/kr/text v0.2.0 // indirect
 	github.com/outcaste-io/ristretto v0.2.1 // indirect
 	github.com/petermattis/goid v0.0.0-20180202154549-b0b1615b78e5 // indirect
 	github.com/philhofer/fwd v1.1.1 // indirect

go.sum

@@ -95,6 +95,7 @@ github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWH
 github.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb/go.mod h1:ZjrT6AXHbDs86ZSdt/osfBi5qfexBrKUdONk989Wnk4=
 github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
 github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -236,6 +237,7 @@ github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfn
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
+github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/mattn/go-zglob v0.0.4 h1:LQi2iOm0/fGgu80AioIJ/1j9w9Oh+9DZ39J4VAGzHQM=
 github.com/mattn/go-zglob v0.0.4/go.mod h1:MxxjyoXXnMxfIpxTK2GAkw1w8glPsQILx3N5wrKakiY=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=