coder · dannykopping · May 30, 2024 · May 15, 2024 · May 15, 2024 · May 15, 2024
diff --git a/provisionersdk/agent.go b/provisionersdk/agent.go
@@ -39,9 +39,9 @@ var (
 	}
 )
 
-// AgentScriptEnv returns a key-pair of scripts that are consumed
-// by the Coder Terraform Provider. See:
-// https://github.com/coder/terraform-provider-coder/blob/main/internal/provider/provider.go#L97
+// AgentScriptEnv returns a key-pair of scripts that are consumed by the Coder Terraform Provider.
+// https://github.com/coder/terraform-provider-coder/blob/main/provider/agent.go (updateInitScript)
+// performs additional string substitutions.
 func AgentScriptEnv() map[string]string {
 	env := map[string]string{}
 	for operatingSystem, scripts := range agentScripts {

diff --git a/provisionersdk/agent_test.go b/provisionersdk/agent_test.go
@@ -7,57 +7,152 @@
 package provisionersdk_test
 
 import (
+	"bytes"
+	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"os/exec"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/go-chi/render"
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/v2/testutil"
+
 	"github.com/coder/coder/v2/provisionersdk"
 )
 
+// mimicking the --version output which we use to test the binary (see provisionersdk/scripts/bootstrap_*).
+const versionOutput = `Coder v2.11.0+8979bfe Tue May  7 17:30:19 UTC 2024`
+
 // bashEcho is a script that calls the local `echo` with the arguments.  This is preferable to
 // sending the real `echo` binary since macOS 14.4+ immediately sigkills `echo` if it is copied to
 // another directory and run locally.
 const bashEcho = `#!/usr/bin/env bash
-echo $@`
+echo "` + versionOutput + `"`
+
+const unexpectedEcho = `#!/usr/bin/env bash
+echo "this is not the agent you are looking for"`
 
 func TestAgentScript(t *testing.T) {
 	t.Parallel()
-	t.Run("Run", func(t *testing.T) {
+
+	t.Run("Valid", func(t *testing.T) {
 		t.Parallel()
-		srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
-			render.Status(r, http.StatusOK)
-			render.Data(rw, r, []byte(bashEcho))
-		}))
-		defer srv.Close()
-		srvURL, err := url.Parse(srv.URL)
-		require.NoError(t, err)
 
-		script, exists := provisionersdk.AgentScriptEnv()[fmt.Sprintf("CODER_AGENT_SCRIPT_%s_%s", runtime.GOOS, runtime.GOARCH)]
-		if !exists {
-			t.Skip("Agent not supported...")
-			return
-		}
-		script = strings.ReplaceAll(script, "${ACCESS_URL}", srvURL.String()+"/")
-		script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token")
+		script := serveScript(t, bashEcho)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		t.Cleanup(cancel)
+
+		var output bytes.Buffer
 		// This is intentionally ran in single quotes to mimic how a customer may
 		// embed our script. Our scripts should not include any single quotes.
 		// nolint:gosec
-		output, err := exec.Command("sh", "-c", "sh -c '"+script+"'").CombinedOutput()
-		t.Log(string(output))
+		cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
+		cmd.Stdout = &output
+		cmd.Stderr = &output
+		require.NoError(t, cmd.Start())
+
+		err := cmd.Wait()
+		if err != nil {
+			var exitErr *exec.ExitError
+			if errors.As(err, &exitErr) {
+				require.Equal(t, 0, exitErr.ExitCode())
+			} else {
+				t.Fatalf("unexpected err: %s", err)
+			}
+		}
+
+		t.Log(output.String())
 		require.NoError(t, err)
 		// Ignore debug output from `set -x`, we're only interested in the last line.
-		lines := strings.Split(strings.TrimSpace(string(output)), "\n")
+		lines := strings.Split(strings.TrimSpace(output.String()), "\n")
 		lastLine := lines[len(lines)-1]
-		// Because we use the "echo" binary, we should expect the arguments provided
+		// When we use the "bashEcho" binary, we should expect the arguments provided
 		// as the response to executing our script.
-		require.Equal(t, "agent", lastLine)
+		require.Equal(t, versionOutput, lastLine)
 	})
+
+	t.Run("Invalid", func(t *testing.T) {
+		t.Parallel()
+
+		script := serveScript(t, unexpectedEcho)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		t.Cleanup(cancel)
+
+		var output bytes.Buffer
+		// This is intentionally ran in single quotes to mimic how a customer may
+		// embed our script. Our scripts should not include any single quotes.
+		// nolint:gosec
+		cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
+		cmd.WaitDelay = time.Second
+		cmd.Stdout = &output
+		cmd.Stderr = &output
+		require.NoError(t, cmd.Start())
+
+		done := make(chan error, 1)
+		var wg sync.WaitGroup
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			// The bootstrap scripts trap exit codes to allow operators to view the script logs and debug the process
+			// while it is still running. We do not expect Wait() to complete.
+			err := cmd.Wait()
+			done <- err
+		}()
+
+		select {
+		case <-ctx.Done():
+			// Timeout.
+			break
+		case err := <-done:
+			// If done signals before context times out, script behaved in an unexpected way.
+			if err != nil {
+				t.Fatalf("unexpected err: %s", err)
+			}
+		}
+
+		// Kill the command, wait for the command to yield.
+		require.NoError(t, cmd.Cancel())
+		wg.Wait()
+
+		t.Log(output.String())
+
+		require.Eventually(t, func() bool {
+			return bytes.Contains(output.Bytes(), []byte("ERROR: Downloaded agent binary returned unexpected version output"))
+		}, testutil.WaitShort, testutil.IntervalSlow)
+	})
+}
+
+// serveScript creates a fake HTTP server which serves a requested "agent binary" (which is actually the given input string)
+// which will be attempted to run to verify that it is correct.
+func serveScript(t *testing.T, in string) string {
+	t.Helper()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		render.Status(r, http.StatusOK)
+		render.Data(rw, r, []byte(in))
+	}))
+	t.Cleanup(srv.Close)
+	srvURL, err := url.Parse(srv.URL)
+	require.NoError(t, err)
+
+	script, exists := provisionersdk.AgentScriptEnv()[fmt.Sprintf("CODER_AGENT_SCRIPT_%s_%s", runtime.GOOS, runtime.GOARCH)]
+	if !exists {
+		t.Skip("Agent not supported...")
+		return ""
+	}
+	script = strings.ReplaceAll(script, "${ACCESS_URL}", srvURL.String()+"/")
+	script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token")
+	return script
 }
diff --git a/provisionersdk/scripts/bootstrap_darwin.sh b/provisionersdk/scripts/bootstrap_darwin.sh
@@ -4,7 +4,7 @@ set -eux
 # This is to allow folks to exec into a failed workspace and poke around to
 # troubleshoot.
 waitonexit() {
-	echo "=== Agent script exited with non-zero code. Sleeping 24h to preserve logs..."
+	echo "=== Agent script exited with non-zero code ($?). Sleeping 24h to preserve logs..."
 	sleep 86400
 }
 trap waitonexit EXIT
@@ -31,4 +31,12 @@ fi
 
 export CODER_AGENT_AUTH="${AUTH_TYPE}"
 export CODER_AGENT_URL="${ACCESS_URL}"
-exec ./$BINARY_NAME agent
+
+output=$(./${BINARY_NAME} --version | head -n1)
+if ! echo "${output}" | grep -q Coder; then
+	echo >&2 "ERROR: Downloaded agent binary returned unexpected version output"
+	echo >&2 "${BINARY_NAME} --version output: \"${output}\""
+	exit 2
+fi
+
+exec ./${BINARY_NAME} agent
diff --git a/provisionersdk/scripts/bootstrap_linux.sh b/provisionersdk/scripts/bootstrap_linux.sh
@@ -4,7 +4,7 @@ set -eux
 # This is to allow folks to exec into a failed workspace and poke around to
 # troubleshoot.
 waitonexit() {
-	echo "=== Agent script exited with non-zero code. Sleeping 24h to preserve logs..."
+	echo "=== Agent script exited with non-zero code ($?). Sleeping 24h to preserve logs..."
 	sleep 86400
 }
 trap waitonexit EXIT
@@ -86,4 +86,12 @@ fi
 
 export CODER_AGENT_AUTH="${AUTH_TYPE}"
 export CODER_AGENT_URL="${ACCESS_URL}"
-exec ./$BINARY_NAME agent
+
+output=$(./${BINARY_NAME} --version | head -n1)
+if ! echo "${output}" | grep -q Coder; then
+	echo >&2 "ERROR: Downloaded agent binary returned unexpected version output"
+	echo >&2 "${BINARY_NAME} --version output: \"${output}\""
+	exit 2
+fi
+
+exec ./${BINARY_NAME} agent
diff --git a/provisionersdk/scripts/bootstrap_windows.ps1 b/provisionersdk/scripts/bootstrap_windows.ps1
@@ -35,6 +35,19 @@ if (-not (Get-Command 'Set-MpPreference' -ErrorAction SilentlyContinue)) {
 $env:CODER_AGENT_AUTH = "${AUTH_TYPE}"
 $env:CODER_AGENT_URL = "${ACCESS_URL}"
 
+$psi = [System.Diagnostics.ProcessStartInfo]::new("$env:TEMP\sshd.exe", '--version')
+$psi.UseShellExecute = $false
+$psi.RedirectStandardOutput = $true
+$p = [System.Diagnostics.Process]::Start($psi)
+$output = $p.StandardOutput.ReadToEnd()
+$p.WaitForExit()
+
+if ($output -notlike "*Coder*") {
+  Write-Output "$env:TEMP\sshd.exe --version output: `"$output"`"
+  Write-Error "ERROR: Downloaded agent binary returned unexpected version output"
+  Throw "unexpected binary"
+}
+
 # Check if we're running inside a Windows container!
 $inContainer = $false
 if ((Get-ItemProperty 'HKLM:\SYSTEM\CurrentControlSet\Control' -Name 'ContainerType' -ErrorAction SilentlyContinue) -ne $null) {

diff --git a/scripts/check_site_icons.sh b/scripts/check_site_icons.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 set -euo pipefail
 # shellcheck source=scripts/lib.sh