agent/agent_test.go
Rama main · réplica de lectura del repositorio original.
package agent_test
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/netip"
"os"
"os/exec"
"os/user"
"path"
"path/filepath"
"regexp"
"runtime"
"slices"
"strconv"
"strings"
"sync"
"testing"
"time"
"github.com/bramvdbogaerde/go-scp"
"github.com/google/uuid"
"github.com/ory/dockertest/v3"
"github.com/ory/dockertest/v3/docker"
"github.com/pion/udp"
"github.com/pkg/sftp"
"github.com/prometheus/client_golang/prometheus"
promgo "github.com/prometheus/client_model/go"
"github.com/spf13/afero"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/goleak"
"golang.org/x/crypto/ssh"
"golang.org/x/xerrors"
"tailscale.com/net/speedtest"
"tailscale.com/tailcfg"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agentcontainers"
"github.com/coder/coder/v2/agent/agentssh"
"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/agent/usershell"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/cryptorand"
"github.com/coder/coder/v2/tailnet"
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
"github.com/coder/coder/v2/tailnet/tailnettest"
"github.com/coder/coder/v2/testutil"
"github.com/coder/coder/v2/testutil/expecter"
"github.com/coder/quartz"
)
func TestMain(m *testing.M) {
if os.Getenv("CODER_TEST_RUN_SUB_AGENT_MAIN") == "1" {
// If we're running as a subagent, we don't want to run the main tests.
// Instead, we just run the subagent tests.
exit := runSubAgentMain()
os.Exit(exit)
}
// Report multi-second pauses of the whole process, see coder/internal#1365.
testutil.StartStallDetector(5 * time.Second)
// Fail with a heap profile before a runaway allocation pushes the runner into swap.
testutil.StartMemoryGuardFromEnv()
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
}
var sshPorts = []uint16{workspacesdk.AgentSSHPort, workspacesdk.AgentStandardSSHPort}
// TestAgent_CloseWhileStarting is a regression test for https://github.com/coder/coder/issues/17328
func TestAgent_ImmediateClose(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
logger := slogtest.Make(t, &slogtest.Options{
// Agent can drop errors when shutting down, and some, like the
// fasthttplistener connection closed error, are unexported.
IgnoreErrors: true,
}).Leveled(slog.LevelDebug)
manifest := agentsdk.Manifest{
AgentID: uuid.New(),
AgentName: "test-agent",
WorkspaceName: "test-workspace",
WorkspaceID: uuid.New(),
}
coordinator := tailnet.NewCoordinator(logger)
t.Cleanup(func() {
_ = coordinator.Close()
})
statsCh := make(chan *proto.Stats, 50)
fs := afero.NewMemMapFs()
client := agenttest.NewClient(t, logger.Named("agenttest"), manifest.AgentID, manifest, statsCh, coordinator)
t.Cleanup(client.Close)
options := agent.Options{
Client: client,
Filesystem: fs,
Logger: logger.Named("agent"),
ReconnectingPTYTimeout: 0,
EnvironmentVariables: map[string]string{},
}
agentUnderTest := agent.New(options)
t.Cleanup(func() {
_ = agentUnderTest.Close()
})
// wait until the agent has connected and is starting to find races in the startup code
_ = testutil.TryReceive(ctx, t, client.GetStartup())
t.Log("Closing Agent")
err := agentUnderTest.Close()
require.NoError(t, err)
}
// eofClient forces the agent's RPC connection to fail with io.EOF so runLoop
// takes the disconnect path.
type eofClient struct {
*agenttest.Client
}
func (*eofClient) ConnectRPC216WithRole(context.Context, string) (
proto.DRPCAgentClient216, tailnetproto.DRPCTailnetClient28, error,
) {
return nil, nil, io.EOF
}
// flushTestSink is a concurrency-safe slog.Sink that records entry messages.
type flushTestSink struct {
mu sync.Mutex
messages []string
}
func (s *flushTestSink) LogEntry(_ context.Context, e slog.SinkEntry) {
s.mu.Lock()
defer s.mu.Unlock()
s.messages = append(s.messages, e.Message)
}
func (*flushTestSink) Sync() {}
func (s *flushTestSink) contains(msg string) bool {
s.mu.Lock()
defer s.mu.Unlock()
return slices.Contains(s.messages, msg)
}
// TestAgent_FlushFlightRecorderOnDisconnect verifies the agent emits its recorded
// debug history when it loses the connection to coderd, so the detail leading
// up to the disconnect is available even though the agent runs at Info.
func TestAgent_FlushFlightRecorderOnDisconnect(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
sink := &flushTestSink{}
logger := slog.Make(sink).Leveled(slog.LevelInfo).FlightRecorder(1000)
// The debug entry is held by the flight recorder, not emitted, during normal
// operation at Info.
const sentinel = "recorded debug sentinel"
logger.Debug(ctx, sentinel)
require.False(t, sink.contains(sentinel), "debug entry should be recorded, not emitted")
manifest := agentsdk.Manifest{
AgentID: uuid.New(),
AgentName: "test-agent",
WorkspaceName: "test-workspace",
WorkspaceID: uuid.New(),
}
coordinator := tailnet.NewCoordinator(logger)
t.Cleanup(func() {
_ = coordinator.Close()
})
statsCh := make(chan *proto.Stats, 50)
baseClient := agenttest.NewClient(t, logger.Named("agenttest"), manifest.AgentID, manifest, statsCh, coordinator)
t.Cleanup(baseClient.Close)
agentUnderTest := agent.New(agent.Options{
Client: &eofClient{Client: baseClient},
Filesystem: afero.NewMemMapFs(),
Logger: logger.Named("agent"),
})
t.Cleanup(func() {
_ = agentUnderTest.Close()
})
// Losing the connection to coderd flushes the recorded debug history to the
// sink.
require.Eventually(t, func() bool {
return sink.contains(sentinel)
}, testutil.WaitShort, testutil.IntervalFast)
}
// NOTE(Cian): I noticed that these tests would fail when my default shell was zsh.
// Writing "exit 0" to stdin before closing fixed the issue for me.
func TestAgent_Stats_SSH(t *testing.T) {
t.Parallel()
for _, port := range sshPorts {
t.Run(fmt.Sprintf("(:%d)", port), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClientOnPort(ctx, port)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
err = session.Shell()
require.NoError(t, err)
// Generate SSH traffic so the connstats window sees the session.
_, err = stdin.Write([]byte("echo test\n"))
require.NoError(t, err)
assertSSHStats(t, stats)
_, err = stdin.Write([]byte("exit 0\n"))
require.NoError(t, err, "writing exit to stdin")
_ = stdin.Close()
err = session.Wait()
require.NoError(t, err, "waiting for session to exit")
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
Reason: string(codersdk.DisconnectReasonGraceful),
RxBytes: 17, // echo + exit
TxBytes: wantSomeBytes, // echo output + prompts
},
)
})
}
// Regression test for CODAGT-517: the barrier blocks reportLoop's
// initial UpdateStats, so on unfixed code the connstats callback is
// never installed and handshake traffic is lost. On fixed code the
// callback is installed at creation, so traffic is captured.
t.Run("StatsCallbackRace", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
barrier := make(chan struct{})
//nolint:dogsled
conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0,
func(c *agenttest.Client, _ *agent.Options) {
c.SetUpdateStatsOverride(func(
ctx context.Context,
req *proto.UpdateStatsRequest,
next func(context.Context, *proto.UpdateStatsRequest) (*proto.UpdateStatsResponse, error),
) (*proto.UpdateStatsResponse, error) {
if req.Stats == nil {
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-barrier:
}
}
return next(ctx, req)
})
},
)
// Connect SSH while the barrier holds reportLoop blocked.
sshClient, err := conn.SSHClientOnPort(ctx, workspacesdk.AgentStandardSSHPort)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
err = session.Shell()
require.NoError(t, err)
// Shell must be idle so the only traffic is the SSH handshake.
close(barrier)
assertSSHStats(t, stats)
_, err = stdin.Write([]byte("exit 0\n"))
require.NoError(t, err, "writing exit to stdin")
_ = stdin.Close()
err = session.Wait()
require.NoError(t, err, "waiting for session to exit")
})
}
// assertSSHStats waits for ConnectionCount, RxBytes, TxBytes, and
// SessionCountSsh to be nonzero on the stats channel.
func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
t.Helper()
var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountSSHSeen bool
require.Eventuallyf(t, func() bool {
s, ok := <-stats
if !ok {
return false
}
t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d",
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["ssh"])
if s.ConnectionCount > 0 {
connectionCountSeen = true
}
if s.RxBytes > 0 {
rxBytesSeen = true
}
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCounts["ssh"] == 1 {
sessionCountSSHSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen
}, testutil.WaitLong, testutil.IntervalFast,
"never saw all SSH stats",
)
}
func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "bash")
require.NoError(t, err)
defer ptyConn.Close()
data, err := json.Marshal(workspacesdk.ReconnectingPTYRequest{
Data: "echo test\r\n",
})
require.NoError(t, err)
_, err = ptyConn.Write(data)
require.NoError(t, err)
var s *proto.Stats
// We are looking for four different stats to be reported. They might not all
// arrive at the same time, so we loop until we've seen them all.
var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountReconnectingPTYSeen bool
require.Eventuallyf(t, func() bool {
var ok bool
s, ok = <-stats
if !ok {
return false
}
if s.ConnectionCount > 0 {
connectionCountSeen = true
}
if s.RxBytes > 0 {
rxBytesSeen = true
}
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCounts["reconnecting_pty"] == 1 {
sessionCountReconnectingPTYSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountReconnectingPTYSeen
}, testutil.WaitLong, testutil.IntervalFast,
"never saw all stats: %+v, saw connectionCount: %t, rxBytes: %t, txBytes: %t, sessionCountReconnectingPTY: %t",
s, connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountReconnectingPTYSeen,
)
}
func TestAgent_Stats_Magic(t *testing.T) {
t.Parallel()
t.Run("StripsEnvironmentVariable", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
session.Setenv(agentssh.AppNameEnvironmentVariable, string(codersdk.AppFamilyVSCode))
defer session.Close()
command := "sh -c 'echo $" + agentssh.AppNameEnvironmentVariable + "'"
expected := ""
if runtime.GOOS == "windows" {
expected = "%" + agentssh.AppNameEnvironmentVariable + "%"
command = "cmd.exe /c echo " + expected
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, expected, strings.TrimSpace(string(output)))
})
t.Run("TracksVSCode", func(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("Sleeping for infinity doesn't work on Windows")
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
session.Setenv(agentssh.AppNameEnvironmentVariable, string(codersdk.AppFamilyVSCode))
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
err = session.Shell()
require.NoError(t, err)
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCounts[vscode]=%d, ConnectionMedianLatencyMS=%f",
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["vscode"], s.ConnectionMedianLatencyMs)
return ok &&
// Ensure that the connection didn't count as a "normal" SSH session.
// This was a special one, so it should be labeled specially in the stats!
s.SessionCounts["vscode"] == 1 &&
// Ensure that connection latency is being counted!
// If it isn't, it's set to -1.
s.ConnectionMedianLatencyMs >= 0
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats",
)
_, err = stdin.Write([]byte("exit 0\n"))
require.NoError(t, err, "writing exit to stdin")
_ = stdin.Close()
err = session.Wait()
require.NoError(t, err)
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH, AppName: string(codersdk.AppFamilyVSCode)},
proto.DisconnectEvent{
Reason: string(codersdk.DisconnectReasonGraceful),
RxBytes: 7, // exit command
TxBytes: wantIgnoreBytes, // might get a prompt
},
)
})
for _, shutdown := range []bool{false, true} {
name := "TracksJetBrains"
if shutdown {
name = "TracksJetBrainsAgentShutdown"
}
t.Run(name, func(t *testing.T) {
t.Parallel()
if runtime.GOOS != "linux" {
t.Skip("JetBrains tracking is only supported on Linux")
}
ctx := testutil.Context(t, testutil.WaitLong)
// JetBrains tracking works by looking at the process name listening on the
// forwarded port. If the process's command line includes the magic string
// we are looking for, then we assume it is a JetBrains editor. So when we
// connect to the port we must ensure the process includes that magic string
// to fool the agent into thinking this is JetBrains. To do this we need to
// spawn an external process (in this case a simple echo server) so we can
// control the process name. The -D here is just to mimic how Java options
// are set but is not necessary as the agent looks only for the magic
// string itself anywhere in the command.
_, b, _, ok := runtime.Caller(0)
require.True(t, ok)
dir := filepath.Join(filepath.Dir(b), "../scripts/echoserver/main.go")
echoServerCmd := exec.Command("go", "run", dir,
"-D", agentssh.MagicProcessCmdlineJetBrains)
stdout, err := echoServerCmd.StdoutPipe()
require.NoError(t, err)
err = echoServerCmd.Start()
require.NoError(t, err)
defer echoServerCmd.Process.Kill()
// The echo server prints its port as the first line.
sc := bufio.NewScanner(stdout)
sc.Scan()
remotePort := sc.Text()
sink := testutil.NewFakeSink(t)
conn, agentClient, stats, _, agnt := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Logger = sink.Logger()
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", remotePort))
require.NoError(t, err)
t.Cleanup(func() {
// always close on failure of test
_ = conn.Close()
_ = tunneledConn.Close()
})
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats with conn open: ok=%t, ConnectionCount=%d, SessionCountJetBrains=%d",
ok, s.ConnectionCount, s.SessionCounts["jetbrains"])
return ok && s.SessionCounts["jetbrains"] == 1
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats with conn open",
)
written := requireEcho(t, tunneledConn)
if shutdown {
require.NoError(t, agnt.Close())
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH, AppName: string(codersdk.AppFamilyJetBrains)},
proto.DisconnectEvent{
Reason: string(codersdk.DisconnectReasonServerShutdown),
RxBytes: int64(written),
TxBytes: int64(written),
},
)
requireDisconnectLog(t, sink, "JetBrains channel closed", "", codersdk.DisconnectReasonServerShutdown)
return
}
// Kill the server and connection after checking for the echo.
_ = echoServerCmd.Process.Kill()
_ = tunneledConn.Close()
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats after disconnect %t, %d",
ok, s.SessionCounts["jetbrains"])
return ok &&
s.SessionCounts["jetbrains"] == 0
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats after conn closes",
)
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH, AppName: string(codersdk.AppFamilyJetBrains)},
proto.DisconnectEvent{
Reason: "normal close",
RxBytes: int64(written),
TxBytes: int64(written),
},
)
})
}
}
func TestAgent_SessionExec(t *testing.T) {
t.Parallel()
for _, port := range sshPorts {
t.Run(fmt.Sprintf("(:%d)", port), func(t *testing.T) {
t.Parallel()
session := setupSSHSessionOnPort(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, port)
command := "echo test"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo test"
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, "test", strings.TrimSpace(string(output)))
})
}
}
func TestAgent_Session_EnvironmentVariables(t *testing.T) {
t.Parallel()
tmpdir := t.TempDir()
// Defined by the coder script runner, hardcoded here since we don't
// have a reference to it.
scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin")
manifest := agentsdk.Manifest{
EnvironmentVariables: map[string]string{
"MY_MANIFEST": "true",
"MY_OVERRIDE": "false",
"MY_SESSION_MANIFEST": "false",
},
}
ctx := testutil.Context(t, testutil.WaitLong)
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ScriptDataDir = tmpdir
opts.EnvironmentVariables["MY_OVERRIDE"] = "true"
opts.EnvInfo = sessionEnvInfo{}
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
t.Cleanup(func() { _ = sshClient.Close() })
for envName, want := range map[string]string{
"CODER": "true", // From the agent.
"MY_MANIFEST": "true", // From the manifest.
"MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest.
"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.
"MY_SESSION": "true", // From the session.
"PATH": scriptBinDir,
} {
t.Run(envName, func(t *testing.T) {
t.Parallel()
got := sessionEnvValue(t, sshClient, envName, map[string]string{
"MY_SESSION": "true",
"MY_SESSION_MANIFEST": "true",
})
if envName == "PATH" {
require.Contains(t, filepath.SplitList(got), want)
return
}
require.Equal(t, want, got)
})
}
}
func TestAgent_Session_SecretInjection(t *testing.T) {
t.Parallel()
manifest := agentsdk.Manifest{
EnvironmentVariables: map[string]string{
"SHOULD_BE_OVERRIDDEN": "manifest-value",
},
}
secrets := []agentsdk.WorkspaceSecret{
{EnvName: "MY_SECRET_ENV", Value: []byte("env-secret-value")},
{FilePath: "/tmp/secret-file", Value: []byte("file-secret-content")},
{EnvName: "BOTH_ENV", FilePath: "/tmp/both-file", Value: []byte("both-value")},
{EnvName: "SHOULD_BE_OVERRIDDEN", Value: []byte("secret-wins")},
}
ctx := testutil.Context(t, testutil.WaitLong)
//nolint:dogsled
conn, _, _, fs, _ := setupAgentWithSecrets(t, manifest, secrets, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.EnvInfo = sessionEnvInfo{}
})
// Verify file injection via the agent's filesystem.
content, err := afero.ReadFile(fs, "/tmp/secret-file")
require.NoError(t, err)
require.Equal(t, "file-secret-content", string(content))
content, err = afero.ReadFile(fs, "/tmp/both-file")
require.NoError(t, err)
require.Equal(t, "both-value", string(content))
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
t.Cleanup(func() { _ = sshClient.Close() })
for envName, want := range map[string]string{
"MY_SECRET_ENV": "env-secret-value",
"BOTH_ENV": "both-value",
"SHOULD_BE_OVERRIDDEN": "secret-wins",
} {
t.Run(envName, func(t *testing.T) {
t.Parallel()
require.Equal(t, want, sessionEnvValue(t, sshClient, envName, nil))
})
}
}
// Environment assertions need a shell that does not run host startup files.
// The default shell can run those files before it executes the SSH command,
// even when that command invokes another shell.
type sessionEnvInfo struct {
usershell.SystemEnvInfo
}
func (sessionEnvInfo) Shell(string) (string, error) {
if runtime.GOOS == "windows" {
return exec.LookPath("cmd.exe")
}
return "/bin/sh", nil
}
func (sessionEnvInfo) ModifyCommand(name string, args ...string) (string, []string) {
if runtime.GOOS == "windows" {
// Disable cmd.exe AutoRun commands from the host registry.
args = append([]string{"/d"}, args...)
}
return name, args
}
func sessionEnvValue(t *testing.T, sshClient *ssh.Client, envName string, sessionEnv map[string]string) string {
t.Helper()
ctx := testutil.Context(t, testutil.WaitLong)
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
go func() {
<-ctx.Done()
_ = session.Close()
}()
for name, value := range sessionEnv {
require.NoError(t, session.Setenv(name, value))
}
stderr := &bytes.Buffer{}
session.Stderr = stderr
command := "printf '%s\\n' \"$" + envName + "\""
if runtime.GOOS == "windows" {
command = `echo %` + envName + `%`
}
out, err := session.Output(command)
if err != nil && ctx.Err() != nil {
t.Fatalf("SSH session deadline expired: %v; output: %q, stderr: %q", ctx.Err(), out, stderr.String())
}
require.NoError(t, err, "output: %q, stderr: %q", out, stderr.String())
return strings.TrimSpace(string(out))
}
func TestAgent_StartupScript_SecretInjection(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("startup script test uses sh syntax")
}
tmpDir := t.TempDir()
secretFilePath := filepath.Join(tmpDir, "secret-file")
envProofPath := filepath.Join(tmpDir, "env-proof")
fileProofPath := filepath.Join(tmpDir, "file-proof")
// The startup script reads the secret env var and the secret file,
// writing both to proof files so we can verify they were available
// at script execution time.
script := fmt.Sprintf(
"echo \"$MY_STARTUP_SECRET\" > %s && cat %s > %s",
envProofPath, secretFilePath, fileProofPath,
)
manifest := agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: script,
Timeout: 30 * time.Second,
RunOnStart: true,
}},
}
secrets := []agentsdk.WorkspaceSecret{
{EnvName: "MY_STARTUP_SECRET", Value: []byte("startup-env-value")},
{FilePath: secretFilePath, Value: []byte("startup-file-content")},
}
// Use the real OS filesystem so that both writeSecretFiles and
// the startup script operate on the same filesystem.
//nolint:dogsled
_, client, _, _, _ := setupAgentWithSecrets(t, manifest, secrets, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.Filesystem = afero.NewOsFs()
})
// Wait for the startup script to complete.
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
}, testutil.WaitLong, testutil.IntervalMedium)
require.Contains(t, got, codersdk.WorkspaceAgentLifecycleReady, "agent never reached ready")
// Verify the startup script could read the secret env var.
envProof, err := os.ReadFile(envProofPath)
require.NoError(t, err)
require.Equal(t, "startup-env-value", strings.TrimSpace(string(envProof)))
// Verify the startup script could read the secret file.
fileProof, err := os.ReadFile(fileProofPath)
require.NoError(t, err)
require.Equal(t, "startup-file-content", string(fileProof))
}
func TestAgent_GitSSH(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $GIT_SSH_COMMAND'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
}
output, err := session.Output(command)
require.NoError(t, err)
require.True(t, strings.HasSuffix(strings.TrimSpace(string(output)), "gitssh --"))
}
func TestAgent_SessionTTYShell(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
// This might be our implementation, or ConPTY itself.
// It's difficult to find extensive tests for it, so
// it seems like it could be either.
t.Skip("ConPTY appears to be inconsistent on Windows.")
}
for _, port := range sshPorts {
t.Run(fmt.Sprintf("(%d)", port), func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitMedium)
logger := testutil.Logger(t)
session := setupSSHSessionOnPort(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, port)
command := "sh"
if runtime.GOOS == "windows" {
command = "cmd.exe"
}
err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
stdout := expecter.NewAttachedToSSHSession(t, session)
stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session)
err = session.Start(command)
require.NoError(t, err)
_ = stdout.Peek(ctx, 1) // wait for the prompt
stdin.WriteLine("echo test")
stdout.ExpectMatch(ctx, "test")
stdin.WriteLine("exit")
err = session.Wait()
require.NoError(t, err)
})
}
}
func TestAgent_SessionTTYExitCode(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
command := "areallynotrealcommand"
err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
err = session.Start(command)
require.NoError(t, err)
err = session.Wait()
exitErr := &ssh.ExitError{}
require.True(t, xerrors.As(err, &exitErr))
if runtime.GOOS == "windows" {
assert.Equal(t, 1, exitErr.ExitStatus())
} else {
assert.Equal(t, 127, exitErr.ExitStatus())
}
}
func TestAgent_Session_TTY_MOTD(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
// This might be our implementation, or ConPTY itself.
// It's difficult to find extensive tests for it, so
// it seems like it could be either.
t.Skip("ConPTY appears to be inconsistent on Windows.")
}
u, err := user.Current()
require.NoError(t, err, "get current user")
name := filepath.Join(u.HomeDir, "motd")
wantMOTD := "Welcome to your Coder workspace!"
wantServiceBanner := "Service banner text goes here"
tests := []struct {
name string
manifest agentsdk.Manifest
banner codersdk.ServiceBannerConfig
expected []string
unexpected []string
expectedRe *regexp.Regexp
}{
{
name: "WithoutServiceBanner",
manifest: agentsdk.Manifest{MOTDFile: name},
banner: codersdk.ServiceBannerConfig{},
expected: []string{wantMOTD},
unexpected: []string{wantServiceBanner},
},
{
name: "WithServiceBanner",
manifest: agentsdk.Manifest{MOTDFile: name},
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: wantServiceBanner,
},
expected: []string{wantMOTD, wantServiceBanner},
},
{
name: "ServiceBannerDisabled",
manifest: agentsdk.Manifest{MOTDFile: name},
banner: codersdk.ServiceBannerConfig{
Enabled: false,
Message: wantServiceBanner,
},
expected: []string{wantMOTD},
unexpected: []string{wantServiceBanner},
},
{
name: "ServiceBannerOnly",
manifest: agentsdk.Manifest{},
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: wantServiceBanner,
},
expected: []string{wantServiceBanner},
unexpected: []string{wantMOTD},
},
{
name: "None",
manifest: agentsdk.Manifest{},
banner: codersdk.ServiceBannerConfig{},
unexpected: []string{wantServiceBanner, wantMOTD},
},
{
name: "CarriageReturns",
manifest: agentsdk.Manifest{},
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: "service\n\nbanner\nhere",
},
expected: []string{"service\r\n\r\nbanner\r\nhere\r\n\r\n"},
unexpected: []string{},
},
{
name: "Trim",
// Enable motd since it will be printed after the banner,
// this ensures that we can test for an exact mount of
// newlines.
manifest: agentsdk.Manifest{
MOTDFile: name,
},
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: "\n\n\n\n\n\nbanner\n\n\n\n\n\n",
},
expectedRe: regexp.MustCompile(`([^\n\r]|^)banner\r\n\r\n[^\r\n]`),
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, test.manifest, test.banner, func(fs afero.Fs) {
err := fs.MkdirAll(filepath.Dir(name), 0o700)
require.NoError(t, err)
err = afero.WriteFile(fs, name, []byte(wantMOTD), 0o600)
require.NoError(t, err)
})
testSessionOutput(t, session, test.expected, test.unexpected, test.expectedRe)
})
}
}
//nolint:tparallel // Sub tests need to run sequentially.
func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
// This might be our implementation, or ConPTY itself.
// It's difficult to find extensive tests for it, so
// it seems like it could be either.
t.Skip("ConPTY appears to be inconsistent on Windows.")
}
// Only the banner updates dynamically; the MOTD file does not.
wantServiceBanner := "Service banner text goes here"
tests := []struct {
banner codersdk.ServiceBannerConfig
expected []string
unexpected []string
}{
{
banner: codersdk.ServiceBannerConfig{},
expected: []string{},
unexpected: []string{wantServiceBanner},
},
{
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: wantServiceBanner,
},
expected: []string{wantServiceBanner},
},
{
banner: codersdk.ServiceBannerConfig{
Enabled: false,
Message: wantServiceBanner,
},
expected: []string{},
unexpected: []string{wantServiceBanner},
},
{
banner: codersdk.ServiceBannerConfig{
Enabled: true,
Message: wantServiceBanner,
},
expected: []string{wantServiceBanner},
unexpected: []string{},
},
{
banner: codersdk.ServiceBannerConfig{},
unexpected: []string{wantServiceBanner},
},
}
setSBInterval := func(_ *agenttest.Client, opts *agent.Options) {
opts.ServiceBannerRefreshInterval = testutil.IntervalFast
}
//nolint:dogsled // Allow the blank identifiers.
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval)
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:paralleltest // These tests need to swap the banner func.
for _, port := range sshPorts {
sshClient, err := conn.SSHClientOnPort(ctx, port)
require.NoError(t, err)
t.Cleanup(func() {
_ = sshClient.Close()
})
for i, test := range tests {
t.Run(fmt.Sprintf("(:%d)/%d", port, i), func(t *testing.T) {
// Set new banner func and wait for the agent to call it to update the
// banner. We wait for two calls to ensure the value has been stored:
// the second call can only begin after the first iteration of
// fetchServiceBannerLoop completes (call + store), so after
// receiving two signals at least one store has happened.
ready := make(chan struct{}, 2)
client.SetAnnouncementBannersFunc(func() ([]codersdk.BannerConfig, error) {
select {
case ready <- struct{}{}:
default:
}
return []codersdk.BannerConfig{test.banner}, nil
})
testutil.TryReceive(ctx, t, ready)
testutil.TryReceive(ctx, t, ready)
session, err := sshClient.NewSession()
require.NoError(t, err)
t.Cleanup(func() {
_ = session.Close()
})
testSessionOutput(t, session, test.expected, test.unexpected, nil)
})
}
}
}
//nolint:paralleltest // This test sets an environment variable.
func TestAgent_Session_TTY_QuietLogin(t *testing.T) {
if runtime.GOOS == "windows" {
// This might be our implementation, or ConPTY itself.
// It's difficult to find extensive tests for it, so
// it seems like it could be either.
t.Skip("ConPTY appears to be inconsistent on Windows.")
}
wantNotMOTD := "Welcome to your Coder workspace!"
wantServiceBanner := "Service banner text goes here"
u, err := user.Current()
require.NoError(t, err, "get current user")
motdPath := filepath.Join(u.HomeDir, "motd")
hushloginPath := filepath.Join(u.HomeDir, ".hushlogin")
// Neither banner nor MOTD should show if not a login shell.
t.Run("NotLogin", func(t *testing.T) {
session := setupSSHSession(t, agentsdk.Manifest{
MOTDFile: motdPath,
}, codersdk.ServiceBannerConfig{
Enabled: true,
Message: wantServiceBanner,
}, func(fs afero.Fs) {
err := afero.WriteFile(fs, motdPath, []byte(wantNotMOTD), 0o600)
require.NoError(t, err, "write motd file")
})
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
wantEcho := "foobar"
command := "echo " + wantEcho
output, err := session.Output(command)
require.NoError(t, err)
require.Contains(t, string(output), wantEcho, "should show echo")
require.NotContains(t, string(output), wantNotMOTD, "should not show motd")
require.NotContains(t, string(output), wantServiceBanner, "should not show service banner")
})
// Only the MOTD should be silenced when hushlogin is present.
t.Run("Hushlogin", func(t *testing.T) {
// The agent writes banners and the MOTD before starting the
// login shell. The working theory is that, under CI contention,
// the host login shell can stall during startup and fail to
// process the PTY exit command, leaving session.Wait blocked.
// Use a self-terminating test shell so unrelated host-shell
// behavior cannot block the quiet-login assertions.
shellPath := filepath.Join(t.TempDir(), "shell")
//nolint:gosec // Executable test shell with test-controlled content.
err := os.WriteFile(shellPath, []byte("#!/bin/sh\nexit 0\n"), 0o700)
require.NoError(t, err, "write test shell")
session := setupSSHSession(t, agentsdk.Manifest{
MOTDFile: motdPath,
}, codersdk.BannerConfig{
Enabled: true,
Message: wantServiceBanner,
}, func(fs afero.Fs) {
err := afero.WriteFile(fs, motdPath, []byte(wantNotMOTD), 0o600)
require.NoError(t, err, "write motd file")
// Place an empty .hushlogin in the user's home so the agent's
// isQuietLogin lookup succeeds and showMOTD is skipped.
err = afero.WriteFile(fs, hushloginPath, []byte{}, 0o600)
require.NoError(t, err, "write hushlogin file")
}, func(_ *agenttest.Client, opts *agent.Options) {
opts.EnvInfo = shellOverrideEnvInfo{shell: shellPath}
})
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
stdout := testutil.NewWaitBuffer()
session.Stdout = stdout
require.NoError(t, session.Shell())
ctx := testutil.Context(t, testutil.WaitShort)
stopClose := context.AfterFunc(ctx, func() { _ = session.Close() })
err = session.Wait()
stopClose()
require.NoError(t, err)
require.Contains(t, stdout.String(), wantServiceBanner, "should show service banner")
require.NotContains(t, stdout.String(), wantNotMOTD, "should not show motd")
})
}
func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
t.Parallel()
// This test is here to prevent regressions where quickly executing
// commands (with TTY) don't sync their output to the SSH session.
//
// See: https://github.com/coder/coder/issues/6656
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
var stdout bytes.Buffer
// NOTE(mafredri): Increase iterations to increase chance of failure,
// assuming bug is present. Limiting GOMAXPROCS further
// increases the chance of failure.
// Using 1000 iterations is basically a guaranteed failure (but let's
// not increase test times needlessly).
// Limit GOMAXPROCS (e.g. `export GOMAXPROCS=1`) to further increase
// chance of failure. Also -race helps.
for i := 0; i < 5; i++ {
func() {
stdout.Reset()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
session.Stdout = &stdout
err = session.Start("echo wazzup")
require.NoError(t, err)
err = session.Wait()
require.NoError(t, err)
require.Contains(t, stdout.String(), "wazzup", "should output greeting")
}()
}
}
func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
t.Parallel()
// This test is here to prevent regressions where a command (with or
// without) a large amount of output would not be fully copied to the
// SSH session. On unix systems, this was fixed by duplicating the file
// descriptor of the PTY master and using it for copying the output.
//
// See: https://github.com/coder/coder/issues/6656
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
var stdout bytes.Buffer
// NOTE(mafredri): Increase iterations to increase chance of failure,
// assuming bug is present.
// Using 10 iterations is basically a guaranteed failure (but let's
// not increase test times needlessly). Run with -race and do not
// limit parallelism (`export GOMAXPROCS=10`) to increase the chance
// of failure.
for i := 0; i < 1; i++ {
func() {
stdout.Reset()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
session.Stdout = &stdout
want := strings.Repeat("wazzup", 1024+1) // ~6KB, +1 because 1024 is a common buffer size.
err = session.Start("echo " + want)
require.NoError(t, err)
err = session.Wait()
require.NoError(t, err)
require.Contains(t, stdout.String(), want, "should output entire greeting")
}()
}
}
func TestAgent_TCPLocalForwarding(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
rl, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer rl.Close()
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
require.True(t, valid)
remotePort := tcpAddr.Port
go echoOnce(t, rl)
sshClient := setupAgentSSHClient(ctx, t)
conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
require.NoError(t, err)
defer conn.Close()
_ = requireEcho(t, conn)
}
func TestAgent_TCPRemoteForwarding(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
sshClient := setupAgentSSHClient(ctx, t)
localhost := netip.MustParseAddr("127.0.0.1")
var randomPort uint16
var ll net.Listener
var err error
for {
randomPort = testutil.RandomPortNoListen(t)
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
ll, err = sshClient.ListenTCP(addr)
if err != nil {
t.Logf("error remote forwarding: %s", err.Error())
select {
case <-ctx.Done():
t.Fatal("timed out getting random listener")
default:
continue
}
}
break
}
defer ll.Close()
go echoOnce(t, ll)
conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", randomPort))
require.NoError(t, err)
defer conn.Close()
_ = requireEcho(t, conn)
}
func TestAgent_TCPLocalForwardingBlocked(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
rl, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer rl.Close()
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
require.True(t, valid)
remotePort := tcpAddr.Port
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockLocalPortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
_, err = sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
require.ErrorContains(t, err, "administratively prohibited")
}
func TestAgent_TCPRemoteForwardingBlocked(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockReversePortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
localhost := netip.MustParseAddr("127.0.0.1")
randomPort := testutil.RandomPortNoListen(t)
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
_, err = sshClient.ListenTCP(addr)
require.ErrorContains(t, err, "tcpip-forward request denied by peer")
}
func TestAgent_UnixLocalForwardingBlocked(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("unix domain sockets are not fully supported on Windows")
}
ctx := testutil.Context(t, testutil.WaitLong)
tmpdir := testutil.TempDirUnixSocket(t)
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
l, err := net.Listen("unix", remoteSocketPath)
require.NoError(t, err)
defer l.Close()
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockLocalPortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
_, err = sshClient.Dial("unix", remoteSocketPath)
require.ErrorContains(t, err, "administratively prohibited")
}
func TestAgent_UnixRemoteForwardingBlocked(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("unix domain sockets are not fully supported on Windows")
}
ctx := testutil.Context(t, testutil.WaitLong)
tmpdir := testutil.TempDirUnixSocket(t)
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockReversePortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
_, err = sshClient.ListenUnix(remoteSocketPath)
require.ErrorContains(t, err, "streamlocal-forward@openssh.com request denied by peer")
}
// TestAgent_LocalBlockedDoesNotAffectReverse verifies that blocking
// local port forwarding does not prevent reverse port forwarding from
// working. A field-name transposition at any plumbing hop would cause
// both directions to be blocked when only one flag is set.
func TestAgent_LocalBlockedDoesNotAffectReverse(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockLocalPortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
// Reverse forwarding must still work.
localhost := netip.MustParseAddr("127.0.0.1")
var ll net.Listener
for {
randomPort := testutil.RandomPortNoListen(t)
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
ll, err = sshClient.ListenTCP(addr)
if err != nil {
t.Logf("error remote forwarding: %s", err.Error())
select {
case <-ctx.Done():
t.Fatal("timed out getting random listener")
default:
continue
}
}
break
}
_ = ll.Close()
}
// TestAgent_ReverseBlockedDoesNotAffectLocal verifies that blocking
// reverse port forwarding does not prevent local port forwarding from
// working.
func TestAgent_ReverseBlockedDoesNotAffectLocal(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
rl, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer rl.Close()
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
require.True(t, valid)
remotePort := tcpAddr.Port
go echoOnce(t, rl)
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockReversePortForwarding = true
})
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
// Local forwarding must still work.
conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
require.NoError(t, err)
defer conn.Close()
_ = requireEcho(t, conn)
}
func TestAgent_UnixLocalForwarding(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("unix domain sockets are not fully supported on Windows")
}
ctx := testutil.Context(t, testutil.WaitLong)
tmpdir := testutil.TempDirUnixSocket(t)
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
l, err := net.Listen("unix", remoteSocketPath)
require.NoError(t, err)
defer l.Close()
go echoOnce(t, l)
sshClient := setupAgentSSHClient(ctx, t)
conn, err := sshClient.Dial("unix", remoteSocketPath)
require.NoError(t, err)
defer conn.Close()
_, err = conn.Write([]byte("test"))
require.NoError(t, err)
b := make([]byte, 4)
_, err = conn.Read(b)
require.NoError(t, err)
require.Equal(t, "test", string(b))
_ = conn.Close()
}
func TestAgent_UnixRemoteForwarding(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("unix domain sockets are not fully supported on Windows")
}
tmpdir := testutil.TempDirUnixSocket(t)
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
ctx := testutil.Context(t, testutil.WaitLong)
sshClient := setupAgentSSHClient(ctx, t)
l, err := sshClient.ListenUnix(remoteSocketPath)
require.NoError(t, err)
defer l.Close()
go echoOnce(t, l)
conn, err := net.Dial("unix", remoteSocketPath)
require.NoError(t, err)
defer conn.Close()
_ = requireEcho(t, conn)
}
func TestAgent_SFTP(t *testing.T) {
t.Parallel()
t.Run("DefaultWorkingDirectory", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
u, err := user.Current()
require.NoError(t, err, "get current user")
home := u.HomeDir
if runtime.GOOS == "windows" {
home = "/" + strings.ReplaceAll(home, "\\", "/")
}
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
client, err := sftp.NewClient(sshClient)
require.NoError(t, err)
defer client.Close()
wd, err := client.Getwd()
require.NoError(t, err, "get working directory")
require.Equal(t, home, wd, "working directory should be user home")
tempFile := filepath.Join(t.TempDir(), "sftp")
// SFTP only accepts unix-y paths.
remoteFile := filepath.ToSlash(tempFile)
if !path.IsAbs(remoteFile) {
// On Windows, e.g. "/C:/Users/...".
remoteFile = path.Join("/", remoteFile)
}
file, err := client.Create(remoteFile)
require.NoError(t, err)
err = file.Close()
require.NoError(t, err)
_, err = os.Stat(tempFile)
require.NoError(t, err)
// Close the client to trigger disconnect event.
_ = client.Close()
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
RxBytes: wantSomeBytes,
TxBytes: wantSomeBytes,
},
)
})
t.Run("CustomWorkingDirectory", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
// Create a custom directory for the agent to use.
customDir := t.TempDir()
expectedDir := customDir
if runtime.GOOS == "windows" {
expectedDir = "/" + strings.ReplaceAll(customDir, "\\", "/")
}
conn, agentClient, _, fs, _ := setupAgent(t, agentsdk.Manifest{
Directory: customDir,
}, 0)
// The agent stats the working directory against its filesystem, so
// the directory must exist there for it to be honored.
require.NoError(t, fs.MkdirAll(customDir, 0o700))
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
client, err := sftp.NewClient(sshClient)
require.NoError(t, err)
defer client.Close()
wd, err := client.Getwd()
require.NoError(t, err, "get working directory")
require.Equal(t, expectedDir, wd, "working directory should be custom directory")
// Close the client to trigger disconnect event.
_ = client.Close()
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
RxBytes: wantSomeBytes,
TxBytes: wantSomeBytes,
},
)
})
t.Run("MissingWorkingDirectory", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
home, err := os.UserHomeDir()
require.NoError(t, err, "get home dir")
if runtime.GOOS == "windows" {
home = "/" + strings.ReplaceAll(home, "\\", "/")
}
// A configured directory that does not exist on the agent's
// filesystem must fall back to the home directory.
missingDir := filepath.Join(t.TempDir(), "does-not-exist")
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: missingDir,
}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
client, err := sftp.NewClient(sshClient)
require.NoError(t, err)
defer client.Close()
wd, err := client.Getwd()
require.NoError(t, err, "get working directory")
require.Equal(t, home, wd, "working directory should fall back to user home")
})
}
func TestAgent_SCP(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
scpClient, err := scp.NewClientBySSH(sshClient)
require.NoError(t, err)
defer scpClient.Close()
tempFile := filepath.Join(t.TempDir(), "scp")
content := "hello world"
err = scpClient.CopyFile(context.Background(), strings.NewReader(content), tempFile, "0755")
require.NoError(t, err)
_, err = os.Stat(tempFile)
require.NoError(t, err)
// Close the client to trigger disconnect event.
scpClient.Close()
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
Reason: string(codersdk.DisconnectReasonGraceful),
RxBytes: wantSomeBytes,
TxBytes: wantSomeBytes,
},
)
}
func TestAgent_FileTransferBlocked(t *testing.T) {
t.Parallel()
assertFileTransferBlocked := func(t *testing.T, errorMessage string) {
// NOTE: Checking content of the error message is flaky. Most likely there is a race condition, which results
// in stopping the client in different phases, and returning different errors:
// - client read the full error message: File transfer has been disabled.
// - client's stream was terminated before reading the error message: EOF
// - client just read the error code (Windows): Process exited with status 65
isErr := strings.Contains(errorMessage, agentssh.BlockedFileTransferErrorMessage) ||
strings.Contains(errorMessage, "EOF") ||
strings.Contains(errorMessage, "Process exited with status 65")
require.True(t, isErr, "Message: "+errorMessage)
}
t.Run("SFTP", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockFileTransfer = true
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
_, err = sftp.NewClient(sshClient)
require.Error(t, err)
assertFileTransferBlocked(t, err.Error())
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
Code: agentssh.BlockedFileTransferErrorCode,
Reason: "file transfer blocked",
},
)
})
t.Run("SCP with go-scp package", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockFileTransfer = true
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
scpClient, err := scp.NewClientBySSH(sshClient)
require.NoError(t, err)
defer scpClient.Close()
tempFile := filepath.Join(t.TempDir(), "scp")
err = scpClient.CopyFile(context.Background(), strings.NewReader("hello world"), tempFile, "0755")
require.Error(t, err)
assertFileTransferBlocked(t, err.Error())
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
Code: agentssh.BlockedFileTransferErrorCode,
Reason: "file transfer blocked",
TxBytes: wantSomeBytes,
},
)
})
t.Run("Forbidden commands", func(t *testing.T) {
t.Parallel()
for _, c := range agentssh.BlockedFileTransferCommands {
t.Run(c, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.BlockFileTransfer = true
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
stdout, err := session.StdoutPipe()
require.NoError(t, err)
//nolint:govet // we don't need `c := c` in Go 1.22
err = session.Start(c)
require.NoError(t, err)
defer session.Close()
msg, err := io.ReadAll(stdout)
require.NoError(t, err)
assertFileTransferBlocked(t, string(msg))
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
proto.DisconnectEvent{
Code: agentssh.BlockedFileTransferErrorCode,
Reason: "file transfer blocked",
TxBytes: wantSomeBytes,
},
)
})
}
})
}
func TestAgent_EnvironmentVariables(t *testing.T) {
t.Parallel()
key := "EXAMPLE"
value := "value"
session := setupSSHSession(t, agentsdk.Manifest{
EnvironmentVariables: map[string]string{
key: value,
},
}, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $" + key + "'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %" + key + "%"
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, value, strings.TrimSpace(string(output)))
}
func TestAgent_EnvironmentVariableExpansion(t *testing.T) {
t.Parallel()
key := "EXAMPLE"
session := setupSSHSession(t, agentsdk.Manifest{
EnvironmentVariables: map[string]string{
key: "$SOMETHINGNOTSET",
},
}, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $" + key + "'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %" + key + "%"
}
output, err := session.Output(command)
require.NoError(t, err)
expect := ""
if runtime.GOOS == "windows" {
expect = "%EXAMPLE%"
}
// Output should be empty, because the variable is not set!
require.Equal(t, expect, strings.TrimSpace(string(output)))
}
func TestAgent_CoderEnvVars(t *testing.T) {
t.Parallel()
manifest := agentsdk.Manifest{
AgentName: "test-agent",
OwnerName: "test-user",
WorkspaceName: "test-workspace",
WorkspaceID: uuid.New(),
WorkspaceBuildID: uuid.New(),
TemplateID: uuid.New(),
TemplateName: "test-template",
}
for key, want := range map[string]string{
"CODER": "true",
"CODER_WORKSPACE_NAME": manifest.WorkspaceName,
"CODER_WORKSPACE_OWNER_NAME": manifest.OwnerName,
"CODER_WORKSPACE_AGENT_NAME": manifest.AgentName,
"CODER_WORKSPACE_ID": manifest.WorkspaceID.String(),
"CODER_WORKSPACE_BUILD_ID": manifest.WorkspaceBuildID.String(),
"CODER_WORKSPACE_TEMPLATE_ID": manifest.TemplateID.String(),
"CODER_WORKSPACE_TEMPLATE_NAME": manifest.TemplateName,
} {
t.Run(key, func(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, manifest, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $" + key + "'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %" + key + "%"
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, want, strings.TrimSpace(string(output)))
})
}
}
func TestAgent_SSHConnectionEnvVars(t *testing.T) {
t.Parallel()
// Note: the SSH_TTY environment variable should only be set for TTYs.
// For some reason this test produces a TTY locally and a non-TTY in CI
// so we don't test for the absence of SSH_TTY.
for _, key := range []string{"SSH_CONNECTION", "SSH_CLIENT"} {
t.Run(key, func(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $" + key + "'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %" + key + "%"
}
output, err := session.Output(command)
require.NoError(t, err)
require.NotEmpty(t, strings.TrimSpace(string(output)))
})
}
}
func TestAgent_SSHConnectionLoginVars(t *testing.T) {
t.Parallel()
envInfo := usershell.SystemEnvInfo{}
u, err := envInfo.User()
require.NoError(t, err, "get current user")
shell, err := envInfo.Shell(u.Username)
require.NoError(t, err, "get current shell")
tests := []struct {
key string
want string
}{
{
key: "USER",
want: u.Username,
},
{
key: "LOGNAME",
want: u.Username,
},
{
key: "SHELL",
want: shell,
},
}
for _, tt := range tests {
t.Run(tt.key, func(t *testing.T) {
t.Parallel()
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
command := "sh -c 'echo $" + tt.key + "'"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo %" + tt.key + "%"
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, tt.want, strings.TrimSpace(string(output)))
})
}
}
// TestAgent_SSHEnvInfoShell verifies that an agent.Options.EnvInfo whose
// Shell() reports a custom shell is piped through to the SSH session, so the
// session command runs under that shell instead of the host default.
func TestAgent_SSHEnvInfoShell(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("the fake shell is a POSIX script")
}
// A fake shell that ignores its arguments and prints a sentinel. The
// sentinel only appears in the session output if the injected Shell() was
// honored. Otherwise the command's own output ("should-not-run") appears.
const marker = "injected-shell-was-used"
shellPath := filepath.Join(t.TempDir(), "fakeshell")
//nolint:gosec // Executable test shell with test-controlled content.
err := os.WriteFile(shellPath, []byte("#!/bin/sh\necho "+marker+"\n"), 0o700)
require.NoError(t, err)
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, func(_ *agenttest.Client, o *agent.Options) {
o.EnvInfo = shellOverrideEnvInfo{shell: shellPath}
})
output, err := session.Output("echo should-not-run")
require.NoError(t, err)
require.Contains(t, string(output), marker)
require.NotContains(t, string(output), "should-not-run")
}
// shellOverrideEnvInfo is a usershell.EnvInfoer that delegates to the system
// implementation but reports a custom shell.
type shellOverrideEnvInfo struct {
usershell.SystemEnvInfo
shell string
}
func (e shellOverrideEnvInfo) Shell(string) (string, error) { return e.shell, nil }
func TestAgent_Metadata(t *testing.T) {
t.Parallel()
echoHello := "echo 'hello'"
t.Run("Once", func(t *testing.T) {
t.Parallel()
//nolint:dogsled
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
{
Key: "greeting1",
Interval: 0,
Script: echoHello,
},
{
Key: "greeting2",
Interval: 1,
Script: echoHello,
},
},
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = testutil.IntervalFast
})
var gotMd map[string]agentsdk.Metadata
require.Eventually(t, func() bool {
gotMd = client.GetMetadata()
return len(gotMd) == 2
}, testutil.WaitShort, testutil.IntervalFast/2)
collectedAt1 := gotMd["greeting1"].CollectedAt
collectedAt2 := gotMd["greeting2"].CollectedAt
require.Eventually(t, func() bool {
gotMd = client.GetMetadata()
if len(gotMd) != 2 {
panic("unexpected number of metadata")
}
return !gotMd["greeting2"].CollectedAt.Equal(collectedAt2)
}, testutil.WaitShort, testutil.IntervalFast/2)
require.Equal(t, gotMd["greeting1"].CollectedAt, collectedAt1, "metadata should not be collected again")
})
t.Run("Many", func(t *testing.T) {
t.Parallel()
//nolint:dogsled
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
{
Key: "greeting",
Interval: 1,
Timeout: 100,
Script: echoHello,
},
},
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = testutil.IntervalFast
})
var gotMd map[string]agentsdk.Metadata
require.Eventually(t, func() bool {
gotMd = client.GetMetadata()
return len(gotMd) == 1
}, testutil.WaitShort, testutil.IntervalFast/2)
collectedAt1 := gotMd["greeting"].CollectedAt
require.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value))
if !assert.Eventually(t, func() bool {
gotMd = client.GetMetadata()
return gotMd["greeting"].CollectedAt.After(collectedAt1)
}, testutil.WaitShort, testutil.IntervalFast/2) {
t.Fatalf("expected metadata to be collected again")
}
})
}
func TestAgentMetadata_Timing(t *testing.T) {
if runtime.GOOS == "windows" {
// Shell scripting in Windows is a pain, and we have already tested
// that the OS logic works in the simpler tests.
t.SkipNow()
}
testutil.SkipIfNotTiming(t)
t.Parallel()
dir := t.TempDir()
const reportInterval = 2
const intervalUnit = 100 * time.Millisecond
var (
greetingPath = filepath.Join(dir, "greeting")
script = "echo hello | tee -a " + greetingPath
)
//nolint:dogsled
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
{
Key: "greeting",
Interval: reportInterval,
Script: script,
},
{
Key: "bad",
Interval: reportInterval,
Script: "exit 1",
},
},
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = intervalUnit
})
require.Eventually(t, func() bool {
return len(client.GetMetadata()) == 2
}, testutil.WaitShort, testutil.IntervalMedium)
for start := time.Now(); time.Since(start) < testutil.WaitMedium; time.Sleep(testutil.IntervalMedium) {
md := client.GetMetadata()
require.Len(t, md, 2, "got: %+v", md)
require.Equal(t, "hello\n", md["greeting"].Value)
require.Equal(t, "run cmd: exit status 1", md["bad"].Error)
greetingByt, err := os.ReadFile(greetingPath)
require.NoError(t, err)
var (
numGreetings = bytes.Count(greetingByt, []byte("hello"))
idealNumGreetings = time.Since(start) / (reportInterval * intervalUnit)
// We allow a 50% error margin because the report loop may backlog
// in CI and other toasters. In production, there is no hard
// guarantee on timing either, and the frontend gives similar
// wiggle room to the staleness of the value.
upperBound = int(idealNumGreetings) + 1
lowerBound = (int(idealNumGreetings) / 2)
)
if idealNumGreetings < 50 {
// There is an insufficient sample size.
continue
}
t.Logf("numGreetings: %d, idealNumGreetings: %d", numGreetings, idealNumGreetings)
// The report loop may slow down on load, but it should never, ever
// speed up.
if numGreetings > upperBound {
t.Fatalf("too many greetings: %d > %d in %v", numGreetings, upperBound, time.Since(start))
} else if numGreetings < lowerBound {
t.Fatalf("too few greetings: %d < %d", numGreetings, lowerBound)
}
}
}
func TestAgent_OrderedStartupScripts(t *testing.T) {
t.Parallel()
if runtime.GOOS == "windows" {
t.Skip("shell scripts")
}
scripts := func(cloneScript, installScript string) (clone, install codersdk.WorkspaceAgentScript) {
clone = codersdk.WorkspaceAgentScript{
ID: uuid.New(), LogSourceID: uuid.New(), DisplayName: "Clone repo",
ResourceAddress: "coder_script.clone_repo", RunOnStart: true,
Script: cloneScript,
}
install = codersdk.WorkspaceAgentScript{
ID: uuid.New(), LogSourceID: uuid.New(), DisplayName: "Install tools",
ResourceAddress: "coder_script.install_tools", RunOnStart: true,
Script: installScript,
Dependencies: []codersdk.WorkspaceAgentScriptDependency{{
PrerequisiteResourceAddress: "coder_script.clone_repo",
Requirement: codersdk.WorkspaceAgentScriptDependencyRequirementSuccess,
}},
}
return clone, install
}
timingFor := func(client *agenttest.Client, id uuid.UUID) *proto.Timing {
for _, timing := range client.GetTimings() {
if bytes.Equal(timing.ScriptId, id[:]) {
return timing
}
}
return nil
}
// collectLines drains the batches delivered so far and keeps the lines
// written to one script's log.
collectLines := func(logsCh <-chan *proto.BatchCreateLogsRequest, logSourceID uuid.UUID, lines *[]string) {
for {
select {
case batch := <-logsCh:
if !bytes.Equal(batch.LogSourceId, logSourceID[:]) {
continue
}
for _, log := range batch.Logs {
*lines = append(*lines, log.Output)
}
default:
return
}
}
}
startAgent := func(t *testing.T, clone, install codersdk.WorkspaceAgentScript) (*agenttest.Client, <-chan *proto.BatchCreateLogsRequest) {
t.Helper()
logsCh := make(chan *proto.BatchCreateLogsRequest, 256)
//nolint:dogsled // only the client is needed
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{clone, install},
}, 0, func(cl *agenttest.Client, _ *agent.Options) {
cl.SetLogsChannel(logsCh)
})
return client, logsCh
}
waitForLifecycle := func(t *testing.T, client *agenttest.Client, want codersdk.WorkspaceAgentLifecycle) {
t.Helper()
require.Eventually(t, func() bool {
return slices.Contains(client.GetLifecycleStates(), want)
}, testutil.WaitLong, testutil.IntervalMedium)
}
t.Run("WaitsThenRuns", func(t *testing.T) {
t.Parallel()
marker := filepath.Join(t.TempDir(), "cloned")
// install_tools exits 1 if it runs before clone_repo has finished.
clone, install := scripts("sleep 0.5 && touch "+marker, "test -f "+marker)
client, logsCh := startAgent(t, clone, install)
waitForLifecycle(t, client, codersdk.WorkspaceAgentLifecycleReady)
require.FileExists(t, marker)
var installLines []string
collectLines(logsCh, install.LogSourceID, &installLines)
require.Empty(t, installLines)
require.Eventually(t, func() bool {
timing := timingFor(client, install.ID)
return timing != nil && timing.Status == proto.Timing_OK
}, testutil.WaitShort, testutil.IntervalFast)
})
t.Run("SkipsAfterFailure", func(t *testing.T) {
t.Parallel()
marker := filepath.Join(t.TempDir(), "cloned")
clone, install := scripts("exit 128", "touch "+marker)
client, logsCh := startAgent(t, clone, install)
waitForLifecycle(t, client, codersdk.WorkspaceAgentLifecycleStartError)
require.NoFileExists(t, marker)
var installLines []string
require.Eventually(t, func() bool {
collectLines(logsCh, install.LogSourceID, &installLines)
return len(installLines) > 0
}, testutil.WaitShort, testutil.IntervalFast)
require.Equal(t, []string{`Skipping script: dependency "Clone repo" failed.`}, installLines)
require.Eventually(t, func() bool {
timing := timingFor(client, install.ID)
return timing != nil &&
timing.Status == proto.Timing_SKIPPED &&
timing.Start.AsTime().Equal(timing.End.AsTime())
}, testutil.WaitShort, testutil.IntervalFast, "the skipped dependent reports one zero-duration timing")
})
}
func TestAgent_Lifecycle(t *testing.T) {
t.Parallel()
t.Run("StartTimeout", func(t *testing.T) {
t.Parallel()
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "sleep 3",
Timeout: time.Millisecond,
RunOnStart: true,
}},
}, 0)
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleStartTimeout,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return slices.Contains(got, want[len(want)-1])
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got[:len(want)])
})
t.Run("StartError", func(t *testing.T) {
t.Parallel()
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "false",
Timeout: 30 * time.Second,
RunOnStart: true,
}},
}, 0)
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleStartError,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return slices.Contains(got, want[len(want)-1])
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got[:len(want)])
})
t.Run("Ready", func(t *testing.T) {
t.Parallel()
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "echo foo",
Timeout: 30 * time.Second,
RunOnStart: true,
}},
}, 0)
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleReady,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got)
})
t.Run("ShuttingDown", func(t *testing.T) {
t.Parallel()
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "sleep 3",
Timeout: 30 * time.Second,
RunOnStop: true,
}},
}, 0)
assert.Eventually(t, func() bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.WaitShort, testutil.IntervalMedium)
// Start close asynchronously so that we an inspect the state.
done := make(chan struct{})
go func() {
defer close(done)
err := closer.Close()
assert.NoError(t, err)
}()
t.Cleanup(func() {
<-done
})
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleReady,
codersdk.WorkspaceAgentLifecycleShuttingDown,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return slices.Contains(got, want[len(want)-1])
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got[:len(want)])
})
t.Run("ShutdownTimeout", func(t *testing.T) {
t.Parallel()
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "sleep 3",
Timeout: time.Millisecond,
RunOnStop: true,
}},
}, 0)
assert.Eventually(t, func() bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.WaitShort, testutil.IntervalMedium)
// Start close asynchronously so that we an inspect the state.
done := make(chan struct{})
go func() {
defer close(done)
err := closer.Close()
assert.NoError(t, err)
}()
t.Cleanup(func() {
<-done
})
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleReady,
codersdk.WorkspaceAgentLifecycleShuttingDown,
codersdk.WorkspaceAgentLifecycleShutdownTimeout,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return slices.Contains(got, want[len(want)-1])
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got[:len(want)])
})
t.Run("ShutdownError", func(t *testing.T) {
t.Parallel()
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "false",
Timeout: 30 * time.Second,
RunOnStop: true,
}},
}, 0)
assert.Eventually(t, func() bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.WaitShort, testutil.IntervalMedium)
// Start close asynchronously so that we an inspect the state.
done := make(chan struct{})
go func() {
defer close(done)
err := closer.Close()
assert.NoError(t, err)
}()
t.Cleanup(func() {
<-done
})
want := []codersdk.WorkspaceAgentLifecycle{
codersdk.WorkspaceAgentLifecycleStarting,
codersdk.WorkspaceAgentLifecycleReady,
codersdk.WorkspaceAgentLifecycleShuttingDown,
codersdk.WorkspaceAgentLifecycleShutdownError,
}
var got []codersdk.WorkspaceAgentLifecycle
assert.Eventually(t, func() bool {
got = client.GetLifecycleStates()
return slices.Contains(got, want[len(want)-1])
}, testutil.WaitShort, testutil.IntervalMedium)
require.Equal(t, want, got[:len(want)])
})
t.Run("ShutdownScriptOnce", func(t *testing.T) {
t.Parallel()
logger := testutil.Logger(t)
ctx := testutil.Context(t, testutil.WaitMedium)
expected := "this-is-shutdown"
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
statsCh := make(chan *proto.Stats, 50)
client := agenttest.NewClient(t,
logger,
uuid.New(),
agentsdk.Manifest{
DERPMap: derpMap,
Scripts: []codersdk.WorkspaceAgentScript{{
ID: uuid.New(),
LogPath: "coder-startup-script.log",
Script: "echo 1",
RunOnStart: true,
}, {
ID: uuid.New(),
LogPath: "coder-shutdown-script.log",
Script: "echo " + expected,
RunOnStop: true,
}},
},
statsCh,
tailnet.NewCoordinator(logger),
)
defer client.Close()
fs := afero.NewMemMapFs()
agent := agent.New(agent.Options{
Client: client,
Logger: logger.Named("agent"),
Filesystem: fs,
})
// agent.Close() loads the shutdown script from the agent metadata.
// The metadata is populated just before execution of the startup script, so it's mandatory to wait
// until the startup starts.
require.Eventually(t, func() bool {
outputPath := filepath.Join(os.TempDir(), "coder-startup-script.log")
content, err := afero.ReadFile(fs, outputPath)
if err != nil {
t.Logf("read file %q: %s", outputPath, err)
return false
}
return len(content) > 0 // something is in the startup log file
}, testutil.WaitShort, testutil.IntervalMedium)
// In order to avoid shutting down the agent before it is fully started and triggering
// errors, we'll wait until the agent is fully up. It's a bit hokey, but among the last things the agent starts
// is the stats reporting, so getting a stats report is a good indication the agent is fully up.
_ = testutil.TryReceive(ctx, t, statsCh)
err := agent.Close()
require.NoError(t, err, "agent should be closed successfully")
outputPath := filepath.Join(os.TempDir(), "coder-shutdown-script.log")
logFirstRead, err := afero.ReadFile(fs, outputPath)
require.NoError(t, err, "log file should be present")
require.Equal(t, expected, string(bytes.TrimSpace(logFirstRead)))
// Make sure that script can't be executed twice.
err = agent.Close()
require.NoError(t, err, "don't need to close the agent twice, no effect")
logSecondRead, err := afero.ReadFile(fs, outputPath)
require.NoError(t, err, "log file should be present")
require.Equal(t, string(bytes.TrimSpace(logFirstRead)), string(bytes.TrimSpace(logSecondRead)))
})
}
func TestAgent_Startup(t *testing.T) {
t.Parallel()
t.Run("EmptyDirectory", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "",
}, 0)
startup := testutil.TryReceive(ctx, t, client.GetStartup())
require.Equal(t, "", startup.GetExpandedDirectory())
})
t.Run("HomeDirectory", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "~",
}, 0)
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, homeDir, startup.GetExpandedDirectory())
})
t.Run("NotAbsoluteDirectory", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "coder/coder",
}, 0)
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, filepath.Join(homeDir, "coder/coder"), startup.GetExpandedDirectory())
})
t.Run("HomeEnvironmentVariable", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "$HOME",
}, 0)
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, homeDir, startup.GetExpandedDirectory())
})
}
//nolint:paralleltest // This test sets an environment variable.
func TestAgent_ReconnectingPTY(t *testing.T) {
if runtime.GOOS == "windows" {
// This might be our implementation, or ConPTY itself.
// It's difficult to find extensive tests for it, so
// it seems like it could be either.
t.Skip("ConPTY appears to be inconsistent on Windows.")
}
backends := []string{"Buffered", "Screen"}
_, err := exec.LookPath("screen")
hasScreen := err == nil
tmuxPath, err := exec.LookPath("tmux")
hasTmux := err == nil
// Make sure UTF-8 works even with locale variables set to C.
t.Setenv("LANG", "C")
t.Setenv("LC_CTYPE", "C")
t.Setenv("LC_ALL", "")
for _, backendType := range backends {
t.Run(backendType, func(t *testing.T) {
if backendType == "Screen" {
if runtime.GOOS != "linux" {
t.Skipf("`screen` is not supported on %s", runtime.GOOS)
} else if !hasScreen {
t.Skip("`screen` not found")
}
} else if hasScreen && runtime.GOOS == "linux" {
// Set up a PATH that does not have screen in it.
bashPath, err := exec.LookPath("bash")
require.NoError(t, err)
dir, err := os.MkdirTemp("/tmp", "coder-test-reconnecting-pty-PATH")
require.NoError(t, err, "create temp dir for reconnecting pty PATH")
err = os.Symlink(bashPath, filepath.Join(dir, "bash"))
require.NoError(t, err, "symlink bash into reconnecting pty PATH")
t.Setenv("PATH", dir)
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
idConnectionReport := uuid.New()
id := uuid.New()
// Test that the connection is reported when closed. This must be tested
// in the first connection because we care about verifying all of these.
netConn0, err := conn.ReconnectingPTY(ctx, idConnectionReport, 80, 80, "bash --norc")
require.NoError(t, err)
_ = netConn0.Close()
assertConnectionReport(t, agentClient,
proto.ConnectEvent{Method: proto.Connection_METHOD_RECONNECTING_PTY},
proto.DisconnectEvent{
Reason: string(codersdk.DisconnectReasonGraceful),
RxBytes: wantSomeBytes, // should get bytes from the init message
// Depending on close timing, the prompt may not send in time, so skip
// checking the sent bytes.
TxBytes: wantIgnoreBytes,
},
)
// --norc disables executing .bashrc, which is often used to customize the bash prompt
netConn1, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
require.NoError(t, err)
defer netConn1.Close()
tr1 := testutil.NewTerminalReader(t, netConn1)
// A second simultaneous connection.
netConn2, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
require.NoError(t, err)
defer netConn2.Close()
tr2 := testutil.NewTerminalReader(t, netConn2)
matchPrompt := func(line string) bool {
return strings.Contains(line, "$ ") || strings.Contains(line, "# ")
}
matchEchoCommand := func(line string) bool {
return strings.Contains(line, "echo test")
}
matchEchoOutput := func(line string) bool {
return strings.Contains(line, "test") && !strings.Contains(line, "echo")
}
matchExitCommand := func(line string) bool {
return strings.Contains(line, "exit")
}
matchExitOutput := func(line string) bool {
return strings.Contains(line, "exit") || strings.Contains(line, "logout")
}
// Wait for the prompt before writing commands. If the command
// arrives before the prompt is written, screen will sometimes put
// the command output on the same line as the command and the test
// will flake.
require.NoError(t, tr1.ReadUntil(ctx, matchPrompt), "find prompt")
require.NoError(t, tr2.ReadUntil(ctx, matchPrompt), "find prompt")
data, err := json.Marshal(workspacesdk.ReconnectingPTYRequest{
Data: "printf '%s\\n' \"$TERM\"\r",
})
require.NoError(t, err)
_, err = netConn1.Write(data)
require.NoError(t, err)
require.NoError(t, tr1.ReadUntilString(ctx, "xterm-256color"), "find TERM output")
require.NoError(t, tr2.ReadUntilString(ctx, "xterm-256color"), "find TERM output")
require.NoError(t, tr1.ReadUntil(ctx, matchPrompt), "find prompt")
require.NoError(t, tr2.ReadUntil(ctx, matchPrompt), "find prompt")
data, err = json.Marshal(workspacesdk.ReconnectingPTYRequest{
Data: "echo test\r",
})
require.NoError(t, err)
_, err = netConn1.Write(data)
require.NoError(t, err)
// Once for typing the command...
require.NoError(t, tr1.ReadUntil(ctx, matchEchoCommand), "find echo command")
// And another time for the actual output.
require.NoError(t, tr1.ReadUntil(ctx, matchEchoOutput), "find echo output")
// Same for the other connection.
require.NoError(t, tr2.ReadUntil(ctx, matchEchoCommand), "find echo command")
require.NoError(t, tr2.ReadUntil(ctx, matchEchoOutput), "find echo output")
_ = netConn1.Close()
_ = netConn2.Close()
netConn3, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
require.NoError(t, err)
defer netConn3.Close()
tr3 := testutil.NewTerminalReader(t, netConn3)
// Same output again!
require.NoError(t, tr3.ReadUntil(ctx, matchEchoCommand), "find echo command")
require.NoError(t, tr3.ReadUntil(ctx, matchEchoOutput), "find echo output")
// Exit should cause the connection to close.
data, err = json.Marshal(workspacesdk.ReconnectingPTYRequest{
Data: "exit\r",
})
require.NoError(t, err)
_, err = netConn3.Write(data)
require.NoError(t, err)
// Once for the input and again for the output.
require.NoError(t, tr3.ReadUntil(ctx, matchExitCommand), "find exit command")
require.NoError(t, tr3.ReadUntil(ctx, matchExitOutput), "find exit output")
// Wait for the connection to close.
require.ErrorIs(t, tr3.ReadUntil(ctx, nil), io.EOF)
// Try a non-shell command. It should output then immediately exit.
netConn4, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "echo test")
require.NoError(t, err)
defer netConn4.Close()
tr4 := testutil.NewTerminalReader(t, netConn4)
require.NoError(t, tr4.ReadUntil(ctx, matchEchoOutput), "find echo output")
require.ErrorIs(t, tr4.ReadUntil(ctx, nil), io.EOF)
// Ensure that UTF-8 is supported. Avoid the terminal emulator because it
// does not appear to support UTF-8, just make sure the bytes that come
// back have the character in it.
netConn5, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "echo ❯")
require.NoError(t, err)
defer netConn5.Close()
bytes, err := io.ReadAll(netConn5)
require.NoError(t, err)
require.Contains(t, string(bytes), "❯")
if !hasTmux {
t.Log("`tmux` not found, skipping tmux glyph regression")
} else {
glyphs := "⚠╭╮╰╯•›│─█▓░▄❯✔╌"
tmuxSocket := "coder-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
t.Cleanup(func() {
_ = exec.Command(tmuxPath, "-L", tmuxSocket, "kill-server").Run()
})
// Keep the pane alive with a shell builtin until the read loop sees
// the glyphs, otherwise tmux can restore the alternate screen first.
command := fmt.Sprintf(
"%s -L %s new-session %q",
strconv.Quote(tmuxPath),
tmuxSocket,
fmt.Sprintf("printf '%%s\\n' '%s'; read _", glyphs),
)
netConn6, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, command)
require.NoError(t, err)
defer netConn6.Close()
var output strings.Builder
buffer := make([]byte, 1024)
deadline := time.Now().Add(testutil.WaitMedium)
for !strings.Contains(output.String(), glyphs) {
if time.Now().After(deadline) {
require.Contains(t, output.String(), glyphs)
}
require.NoError(t, netConn6.SetReadDeadline(time.Now().Add(testutil.IntervalMedium)))
read, err := netConn6.Read(buffer)
if read > 0 {
_, _ = output.Write(buffer[:read])
}
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
continue
}
require.NoError(t, err)
}
}
})
}
}
// This tests end-to-end functionality of connecting to a running container
// and executing a command. It creates a real Docker container and runs a
// command. As such, it does not run by default in CI.
// You can run it manually as follows:
//
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_ReconnectingPTYContainer
func TestAgent_ReconnectingPTYContainer(t *testing.T) {
t.Parallel()
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
}
if _, err := exec.LookPath("devcontainer"); err != nil {
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
}
pool, err := dockertest.NewPool("")
require.NoError(t, err, "Could not connect to docker")
ct, err := pool.RunWithOptions(&dockertest.RunOptions{
Repository: "busybox",
Tag: "latest",
Cmd: []string{"sleep", "infnity"},
}, func(config *docker.HostConfig) {
config.AutoRemove = true
config.RestartPolicy = docker.RestartPolicy{Name: "no"}
})
require.NoError(t, err, "Could not start container")
defer func() {
err := pool.Purge(ct)
require.NoError(t, err, "Could not stop container")
}()
// Wait for container to start
require.Eventually(t, func() bool {
ct, ok := pool.ContainerByName(ct.Container.Name)
return ok && ct.Container.State.Running
}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
// nolint: dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
)
})
ctx := testutil.Context(t, testutil.WaitLong)
ac, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "/bin/sh", func(arp *workspacesdk.AgentReconnectingPTYInit) {
arp.Container = ct.Container.ID
})
require.NoError(t, err, "failed to create ReconnectingPTY")
defer ac.Close()
tr := testutil.NewTerminalReader(t, ac)
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
return strings.Contains(line, "#") || strings.Contains(line, "$")
}), "find prompt")
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
Data: "hostname\r",
}), "write hostname")
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
return strings.Contains(line, "hostname")
}), "find hostname command")
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
return strings.Contains(line, ct.Container.Config.Hostname)
}), "find hostname output")
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
Data: "exit\r",
}), "write exit command")
// Wait for the connection to close.
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
}
type subAgentRequestPayload struct {
Token string `json:"token"`
Directory string `json:"directory"`
}
// runSubAgentMain is the main function for the sub-agent that connects
// to the control plane. It reads the CODER_AGENT_URL and
// CODER_AGENT_TOKEN environment variables, sends the token, and exits
// with a status code based on the response.
func runSubAgentMain() int {
url := os.Getenv("CODER_AGENT_URL")
token := os.Getenv("CODER_AGENT_TOKEN")
if url == "" || token == "" {
_, _ = fmt.Fprintln(os.Stderr, "CODER_AGENT_URL and CODER_AGENT_TOKEN must be set")
return 10
}
dir, err := os.Getwd()
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "failed to get current working directory: %v\n", err)
return 1
}
payload := subAgentRequestPayload{
Token: token,
Directory: dir,
}
b, err := json.Marshal(payload)
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "failed to marshal payload: %v\n", err)
return 1
}
req, err := http.NewRequest("POST", url, bytes.NewReader(b))
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "failed to create request: %v\n", err)
return 1
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
req = req.WithContext(ctx)
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "agent connection failed: %v\n", err)
return 11
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
_, _ = fmt.Fprintf(os.Stderr, "agent exiting with non-zero exit code %d\n", resp.StatusCode)
return 12
}
_, _ = fmt.Println("sub-agent connected successfully")
return 0
}
// This tests end-to-end functionality of auto-starting a devcontainer.
// It runs "devcontainer up" which creates a real Docker container. As
// such, it does not run by default in CI.
//
// You can run it manually as follows:
//
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerAutostart
//
//nolint:paralleltest // This test sets an environment variable.
func TestAgent_DevcontainerAutostart(t *testing.T) {
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
}
if _, err := exec.LookPath("devcontainer"); err != nil {
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
}
// This HTTP handler handles requests from runSubAgentMain which
// acts as a fake sub-agent. We want to verify that the sub-agent
// connects and sends its token. We use a channel to signal
// that the sub-agent has connected successfully and then we wait
// until we receive another signal to return from the handler. This
// keeps the agent "alive" for as long as we want.
subAgentConnected := make(chan subAgentRequestPayload, 1)
subAgentReady := make(chan struct{}, 1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/api/v2/workspaceagents/me/") {
return
}
t.Logf("Sub-agent request received: %s %s", r.Method, r.URL.Path)
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
// Read the token from the request body.
var payload subAgentRequestPayload
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
http.Error(w, "Failed to read token", http.StatusBadRequest)
t.Logf("Failed to read token: %v", err)
return
}
defer r.Body.Close()
t.Logf("Sub-agent request payload received: %+v", payload)
// Signal that the sub-agent has connected successfully.
select {
case <-t.Context().Done():
t.Logf("Test context done, not processing sub-agent request")
return
case subAgentConnected <- payload:
}
// Wait for the signal to return from the handler.
select {
case <-t.Context().Done():
t.Logf("Test context done, not waiting for sub-agent ready")
return
case <-subAgentReady:
}
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
pool, err := dockertest.NewPool("")
require.NoError(t, err, "Could not connect to docker")
// Prepare temporary devcontainer for test (mywork).
devcontainerID := uuid.New()
tmpdir := t.TempDir()
t.Setenv("HOME", tmpdir)
tempWorkspaceFolder := filepath.Join(tmpdir, "mywork")
unexpandedWorkspaceFolder := filepath.Join("~", "mywork")
t.Logf("Workspace folder: %s", tempWorkspaceFolder)
t.Logf("Unexpanded workspace folder: %s", unexpandedWorkspaceFolder)
devcontainerPath := filepath.Join(tempWorkspaceFolder, ".devcontainer")
err = os.MkdirAll(devcontainerPath, 0o755)
require.NoError(t, err, "create devcontainer directory")
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
err = os.WriteFile(devcontainerFile, []byte(`{
"name": "mywork",
"image": "ubuntu:latest",
"cmd": ["sleep", "infinity"],
"runArgs": ["--network=host", "--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"]
}`), 0o600)
require.NoError(t, err, "write devcontainer.json")
manifest := agentsdk.Manifest{
// Set up pre-conditions for auto-starting a devcontainer, the script
// is expected to be prepared by the provisioner normally.
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
{
ID: devcontainerID,
Name: "test",
// Use an unexpanded path to test the expansion.
WorkspaceFolder: unexpandedWorkspaceFolder,
},
},
Scripts: []codersdk.WorkspaceAgentScript{
{
ID: devcontainerID,
LogSourceID: agentsdk.ExternalLogSourceID,
RunOnStart: true,
Script: "echo this-will-be-replaced",
DisplayName: "Dev Container (test)",
},
},
}
mClock := quartz.NewMock(t)
mClock.Set(time.Now())
tickerFuncTrap := mClock.Trap().TickerFunc("agentcontainers")
//nolint:dogsled
_, agentClient, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
o.DevcontainerAPIOptions = append(
o.DevcontainerAPIOptions,
// Only match this specific dev container.
agentcontainers.WithClock(mClock),
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", tempWorkspaceFolder),
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
agentcontainers.WithSubAgentURL(srv.URL),
// The agent will copy "itself", but in the case of this test, the
// agent is actually this test binary. So we'll tell the test binary
// to execute the sub-agent main function via this env.
agentcontainers.WithSubAgentEnv("CODER_TEST_RUN_SUB_AGENT_MAIN=1"),
)
})
t.Logf("Waiting for container with label: devcontainer.local_folder=%s", tempWorkspaceFolder)
var container docker.APIContainers
require.Eventually(t, func() bool {
containers, err := pool.Client.ListContainers(docker.ListContainersOptions{All: true})
if err != nil {
t.Logf("Error listing containers: %v", err)
return false
}
for _, c := range containers {
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
if labelValue, ok := c.Labels["devcontainer.local_folder"]; ok {
if labelValue == tempWorkspaceFolder {
t.Logf("Found matching container: %s", c.ID[:12])
container = c
return true
}
}
}
return false
}, testutil.WaitSuperLong, testutil.IntervalMedium, "no container with workspace folder label found")
defer func() {
// We can't rely on pool here because the container is not
// managed by it (it is managed by @devcontainer/cli).
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
ID: container.ID,
RemoveVolumes: true,
Force: true,
})
assert.NoError(t, err, "remove container")
}()
containerInfo, err := pool.Client.InspectContainer(container.ID)
require.NoError(t, err, "inspect container")
t.Logf("Container state: status: %v", containerInfo.State.Status)
require.True(t, containerInfo.State.Running, "container should be running")
ctx := testutil.Context(t, testutil.WaitLong)
// Ensure the container update routine runs.
tickerFuncTrap.MustWait(ctx).MustRelease(ctx)
tickerFuncTrap.Close()
// Since the agent does RefreshContainers, and the ticker function
// is set to skip instead of queue, we must advance the clock
// multiple times to ensure that the sub-agent is created.
var subAgents []*proto.SubAgent
for {
_, next := mClock.AdvanceNext()
next.MustWait(ctx)
// Verify that a subagent was created.
subAgents = agentClient.GetSubAgents()
if len(subAgents) > 0 {
t.Logf("Found sub-agents: %d", len(subAgents))
break
}
}
require.Len(t, subAgents, 1, "expected one sub agent")
subAgent := subAgents[0]
subAgentID, err := uuid.FromBytes(subAgent.GetId())
require.NoError(t, err, "failed to parse sub-agent ID")
t.Logf("Connecting to sub-agent: %s (ID: %s)", subAgent.Name, subAgentID)
gotDir, err := agentClient.GetSubAgentDirectory(subAgentID)
require.NoError(t, err, "failed to get sub-agent directory")
require.Equal(t, "/workspaces/mywork", gotDir, "sub-agent directory should match")
subAgentToken, err := uuid.FromBytes(subAgent.GetAuthToken())
require.NoError(t, err, "failed to parse sub-agent token")
payload := testutil.RequireReceive(ctx, t, subAgentConnected)
require.Equal(t, subAgentToken.String(), payload.Token, "sub-agent token should match")
require.Equal(t, "/workspaces/mywork", payload.Directory, "sub-agent directory should match")
// Allow the subagent to exit.
close(subAgentReady)
}
// TestAgent_DevcontainerRecreate tests that RecreateDevcontainer
// recreates a devcontainer and emits logs.
//
// This tests end-to-end functionality of auto-starting a devcontainer.
// It runs "devcontainer up" which creates a real Docker container. As
// such, it does not run by default in CI.
//
// You can run it manually as follows:
//
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerRecreate
func TestAgent_DevcontainerRecreate(t *testing.T) {
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
}
t.Parallel()
pool, err := dockertest.NewPool("")
require.NoError(t, err, "Could not connect to docker")
// Prepare temporary devcontainer for test (mywork).
devcontainerID := uuid.New()
devcontainerLogSourceID := uuid.New()
workspaceFolder := filepath.Join(t.TempDir(), "mywork")
t.Logf("Workspace folder: %s", workspaceFolder)
devcontainerPath := filepath.Join(workspaceFolder, ".devcontainer")
err = os.MkdirAll(devcontainerPath, 0o755)
require.NoError(t, err, "create devcontainer directory")
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
err = os.WriteFile(devcontainerFile, []byte(`{
"name": "mywork",
"image": "busybox:latest",
"cmd": ["sleep", "infinity"],
"runArgs": ["--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"]
}`), 0o600)
require.NoError(t, err, "write devcontainer.json")
manifest := agentsdk.Manifest{
// Set up pre-conditions for auto-starting a devcontainer, the
// script is used to extract the log source ID.
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
{
ID: devcontainerID,
Name: "test",
WorkspaceFolder: workspaceFolder,
},
},
Scripts: []codersdk.WorkspaceAgentScript{
{
ID: devcontainerID,
LogSourceID: devcontainerLogSourceID,
},
},
}
//nolint:dogsled
conn, client, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", workspaceFolder),
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
)
})
ctx := testutil.Context(t, testutil.WaitLong)
// We enabled autostart for the devcontainer, so ready is a good
// indication that the devcontainer is up and running. Importantly,
// this also means that the devcontainer startup is no longer
// producing logs that may interfere with the recreate logs.
testutil.Eventually(ctx, t, func(context.Context) bool {
states := client.GetLifecycleStates()
return slices.Contains(states, codersdk.WorkspaceAgentLifecycleReady)
}, testutil.IntervalMedium, "devcontainer not ready")
t.Logf("Looking for container with label: devcontainer.local_folder=%s", workspaceFolder)
var container codersdk.WorkspaceAgentContainer
testutil.Eventually(ctx, t, func(context.Context) bool {
resp, err := conn.ListContainers(ctx)
if err != nil {
t.Logf("Error listing containers: %v", err)
return false
}
for _, c := range resp.Containers {
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
if v, ok := c.Labels["devcontainer.local_folder"]; ok && v == workspaceFolder {
t.Logf("Found matching container: %s", c.ID[:12])
container = c
return true
}
}
return false
}, testutil.IntervalMedium, "no container with workspace folder label found")
defer func(container codersdk.WorkspaceAgentContainer) {
// We can't rely on pool here because the container is not
// managed by it (it is managed by @devcontainer/cli).
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
ID: container.ID,
RemoveVolumes: true,
Force: true,
})
assert.Error(t, err, "container should be removed by recreate")
}(container)
ctx = testutil.Context(t, testutil.WaitLong) // Reset context.
// Capture logs via ScriptLogger.
logsCh := make(chan *proto.BatchCreateLogsRequest, 1)
client.SetLogsChannel(logsCh)
// Invoke recreate to trigger the destruction and recreation of the
// devcontainer, we do it in a goroutine so we can process logs
// concurrently.
go func(container codersdk.WorkspaceAgentContainer) {
_, err := conn.RecreateDevcontainer(ctx, devcontainerID.String())
assert.NoError(t, err, "recreate devcontainer should succeed")
}(container)
t.Logf("Checking recreate logs for outcome...")
// Wait for the logs to be emitted, the @devcontainer/cli up command
// will emit a log with the outcome at the end suggesting we did
// receive all the logs.
waitForOutcomeLoop:
for {
batch := testutil.RequireReceive(ctx, t, logsCh)
if bytes.Equal(batch.LogSourceId, devcontainerLogSourceID[:]) {
for _, log := range batch.Logs {
t.Logf("Received log: %s", log.Output)
if strings.Contains(log.Output, "\"outcome\"") {
break waitForOutcomeLoop
}
}
}
}
t.Logf("Checking there's a new container with label: devcontainer.local_folder=%s", workspaceFolder)
// Make sure the container exists and isn't the same as the old one.
testutil.Eventually(ctx, t, func(context.Context) bool {
resp, err := conn.ListContainers(ctx)
if err != nil {
t.Logf("Error listing containers: %v", err)
return false
}
for _, c := range resp.Containers {
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
if v, ok := c.Labels["devcontainer.local_folder"]; ok && v == workspaceFolder {
if c.ID == container.ID {
t.Logf("Found same container: %s", c.ID[:12])
return false
}
t.Logf("Found new container: %s", c.ID[:12])
container = c
return true
}
}
return false
}, testutil.IntervalMedium, "new devcontainer not found")
defer func(container codersdk.WorkspaceAgentContainer) {
// We can't rely on pool here because the container is not
// managed by it (it is managed by @devcontainer/cli).
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
ID: container.ID,
RemoveVolumes: true,
Force: true,
})
assert.NoError(t, err, "remove container")
}(container)
}
func TestAgent_DevcontainersDisabledForSubAgent(t *testing.T) {
t.Parallel()
// Create a manifest with a ParentID to make this a sub agent.
manifest := agentsdk.Manifest{
AgentID: uuid.New(),
ParentID: uuid.New(),
}
// Setup the agent with devcontainers enabled initially.
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
})
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
defer cancel()
var err error
// setupAgent only waits for tailnet reachability, not for the HTTP API
// listener to serve the expected sub-agent rejection response.
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
_, err = conn.ListContainers(ctx)
if err != nil {
t.Logf("Error listing containers: %v", err)
}
return err != nil && strings.Contains(err.Error(), "Dev Container feature not supported.")
}, testutil.IntervalFast, "containers endpoint should reject devcontainers inside sub agents")
require.Error(t, err)
require.Contains(t, err.Error(), "Dev Container feature not supported.")
require.Contains(t, err.Error(), "Dev Container integration inside other Dev Containers is explicitly not supported.")
}
// TestAgent_DevcontainerPrebuildClaim tests that we correctly handle
// the claiming process for running devcontainers.
//
// You can run it manually as follows:
//
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerPrebuildClaim
//
//nolint:paralleltest // This test sets an environment variable.
func TestAgent_DevcontainerPrebuildClaim(t *testing.T) {
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
}
if _, err := exec.LookPath("devcontainer"); err != nil {
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
}
pool, err := dockertest.NewPool("")
require.NoError(t, err, "Could not connect to docker")
var (
ctx = testutil.Context(t, testutil.WaitShort)
devcontainerID = uuid.New()
devcontainerLogSourceID = uuid.New()
workspaceFolder = filepath.Join(t.TempDir(), "project")
devcontainerPath = filepath.Join(workspaceFolder, ".devcontainer")
devcontainerConfig = filepath.Join(devcontainerPath, "devcontainer.json")
)
// Given: A devcontainer project.
t.Logf("Workspace folder: %s", workspaceFolder)
err = os.MkdirAll(devcontainerPath, 0o755)
require.NoError(t, err, "create dev container directory")
// Given: This devcontainer project specifies an app that uses the owner name and workspace name.
err = os.WriteFile(devcontainerConfig, []byte(`{
"name": "project",
"image": "busybox:latest",
"cmd": ["sleep", "infinity"],
"runArgs": ["--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"],
"customizations": {
"coder": {
"apps": [{
"slug": "zed",
"url": "zed://ssh/${localEnv:CODER_WORKSPACE_AGENT_NAME}.${localEnv:CODER_WORKSPACE_NAME}.${localEnv:CODER_WORKSPACE_OWNER_NAME}.coder${containerWorkspaceFolder}"
}]
}
}
}`), 0o600)
require.NoError(t, err, "write devcontainer config")
// Given: A manifest with a prebuild username and workspace name.
manifest := agentsdk.Manifest{
OwnerName: "prebuilds",
WorkspaceName: "prebuilds-xyz-123",
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
{ID: devcontainerID, Name: "test", WorkspaceFolder: workspaceFolder},
},
Scripts: []codersdk.WorkspaceAgentScript{
{ID: devcontainerID, LogSourceID: devcontainerLogSourceID},
},
}
// When: We create an agent with devcontainers enabled.
//nolint:dogsled
conn, client, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerLocalFolderLabel, workspaceFolder),
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
)
})
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.IntervalMedium, "agent not ready")
var dcPrebuild codersdk.WorkspaceAgentDevcontainer
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
resp, err := conn.ListContainers(ctx)
require.NoError(t, err)
for _, dc := range resp.Devcontainers {
if dc.Container == nil {
continue
}
v, ok := dc.Container.Labels[agentcontainers.DevcontainerLocalFolderLabel]
if ok && v == workspaceFolder {
dcPrebuild = dc
return true
}
}
return false
}, testutil.IntervalMedium, "devcontainer not found")
defer func() {
pool.Client.RemoveContainer(docker.RemoveContainerOptions{
ID: dcPrebuild.Container.ID,
RemoveVolumes: true,
Force: true,
})
}()
// Then: We expect a sub agent to have been created.
subAgents := client.GetSubAgents()
require.Len(t, subAgents, 1)
subAgent := subAgents[0]
subAgentID, err := uuid.FromBytes(subAgent.GetId())
require.NoError(t, err)
// And: We expect there to be 1 app.
subAgentApps, err := client.GetSubAgentApps(subAgentID)
require.NoError(t, err)
require.Len(t, subAgentApps, 1)
// And: This app should contain the prebuild workspace name and owner name.
subAgentApp := subAgentApps[0]
require.Equal(t, "zed://ssh/project.prebuilds-xyz-123.prebuilds.coder/workspaces/project", subAgentApp.GetUrl())
// Given: We close the client and connection
client.Close()
conn.Close()
// Given: A new manifest with a regular user owner name and workspace name.
manifest = agentsdk.Manifest{
OwnerName: "user",
WorkspaceName: "user-workspace",
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
{ID: devcontainerID, Name: "test", WorkspaceFolder: workspaceFolder},
},
Scripts: []codersdk.WorkspaceAgentScript{
{ID: devcontainerID, LogSourceID: devcontainerLogSourceID},
},
}
// When: We create an agent with devcontainers enabled.
//nolint:dogsled
conn, client, _, _, _ = setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
o.Devcontainers = true
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerLocalFolderLabel, workspaceFolder),
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
)
})
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.IntervalMedium, "agent not ready")
var dcClaimed codersdk.WorkspaceAgentDevcontainer
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
resp, err := conn.ListContainers(ctx)
require.NoError(t, err)
for _, dc := range resp.Devcontainers {
if dc.Container == nil {
continue
}
v, ok := dc.Container.Labels[agentcontainers.DevcontainerLocalFolderLabel]
if ok && v == workspaceFolder {
dcClaimed = dc
return true
}
}
return false
}, testutil.IntervalMedium, "devcontainer not found")
defer func() {
if dcClaimed.Container.ID != dcPrebuild.Container.ID {
pool.Client.RemoveContainer(docker.RemoveContainerOptions{
ID: dcClaimed.Container.ID,
RemoveVolumes: true,
Force: true,
})
}
}()
// Then: We expect the claimed devcontainer and prebuild devcontainer
// to be using the same underlying container.
require.Equal(t, dcPrebuild.Container.ID, dcClaimed.Container.ID)
// And: We expect there to be a sub agent created.
subAgents = client.GetSubAgents()
require.Len(t, subAgents, 1)
subAgent = subAgents[0]
subAgentID, err = uuid.FromBytes(subAgent.GetId())
require.NoError(t, err)
// And: We expect there to be an app.
subAgentApps, err = client.GetSubAgentApps(subAgentID)
require.NoError(t, err)
require.Len(t, subAgentApps, 1)
// And: We expect this app to have the user's owner name and workspace name.
subAgentApp = subAgentApps[0]
require.Equal(t, "zed://ssh/project.user-workspace.user.coder/workspaces/project", subAgentApp.GetUrl())
}
func TestAgent_Dial(t *testing.T) {
t.Parallel()
cases := []struct {
name string
setup func(t testing.TB) net.Listener
}{
{
name: "TCP",
setup: func(t testing.TB) net.Listener {
l, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err, "create TCP listener")
return l
},
},
{
name: "UDP",
setup: func(t testing.TB) net.Listener {
addr := net.UDPAddr{
IP: net.ParseIP("127.0.0.1"),
Port: 0,
}
l, err := udp.Listen("udp", &addr)
require.NoError(t, err, "create UDP listener")
return l
},
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
t.Parallel()
// The purpose of this test is to ensure that a client can dial a
// listener in the workspace over tailnet.
//
// The OS sometimes drops packets if the system can't keep up with
// them. For TCP packets, it's typically fine due to
// retransmissions, but for UDP packets, it can fail this test.
//
// The OS gets involved for the Wireguard traffic (either via DERP
// or direct UDP), and also for the traffic between the agent and
// the listener in the "workspace".
//
// To avoid this, we'll retry this test up to 3 times.
//nolint:gocritic // This test is flaky due to uncontrollable OS packet drops under heavy load.
testutil.RunRetry(t, 3, func(t testing.TB) {
ctx := testutil.Context(t, testutil.WaitLong)
l := c.setup(t)
done := make(chan struct{})
defer func() {
l.Close()
<-done
}()
go func() {
defer close(done)
for range 2 {
c, err := l.Accept()
if assert.NoError(t, err, "accept connection") {
testAccept(ctx, t, c)
_ = c.Close()
}
}
}()
agentID := uuid.UUID{0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8}
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
AgentID: agentID,
}, 0)
require.True(t, agentConn.AwaitReachable(ctx))
conn, err := agentConn.DialContext(ctx, l.Addr().Network(), l.Addr().String())
require.NoError(t, err)
testDial(ctx, t, conn)
err = conn.Close()
require.NoError(t, err)
// also connect via the CoderServicePrefix, to test that we can reach the agent on this
// IP. This will be required for CoderVPN.
_, rawPort, _ := net.SplitHostPort(l.Addr().String())
port, _ := strconv.ParseUint(rawPort, 10, 16)
ipp := netip.AddrPortFrom(tailnet.CoderServicePrefix.AddrFromUUID(agentID), uint16(port))
switch l.Addr().Network() {
case "tcp":
conn, err = agentConn.TailnetConn().DialContextTCP(ctx, ipp)
case "udp":
conn, err = agentConn.TailnetConn().DialContextUDP(ctx, ipp)
default:
t.Fatalf("unknown network: %s", l.Addr().Network())
}
require.NoError(t, err)
testDial(ctx, t, conn)
err = conn.Close()
require.NoError(t, err)
})
})
}
}
// TestAgent_UpdatedDERP checks that agents can handle their DERP map being
// updated, and that clients can also handle it.
func TestAgent_UpdatedDERP(t *testing.T) {
t.Parallel()
logger := testutil.Logger(t)
originalDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
require.NotNil(t, originalDerpMap)
coordinator := tailnet.NewCoordinator(logger)
// use t.Cleanup so the coordinator closing doesn't deadlock with in-memory
// coordination
t.Cleanup(func() {
_ = coordinator.Close()
})
agentID := uuid.New()
statsCh := make(chan *proto.Stats, 50)
fs := afero.NewMemMapFs()
client := agenttest.NewClient(t,
logger.Named("agent"),
agentID,
agentsdk.Manifest{
DERPMap: originalDerpMap,
// Force DERP.
DisableDirectConnections: true,
},
statsCh,
coordinator,
)
t.Cleanup(func() {
t.Log("closing client")
client.Close()
})
uut := agent.New(agent.Options{
Client: client,
Filesystem: fs,
Logger: logger.Named("agent"),
ReconnectingPTYTimeout: time.Minute,
})
t.Cleanup(func() {
t.Log("closing agent")
_ = uut.Close()
})
// Setup a client connection.
newClientConn := func(derpMap *tailcfg.DERPMap, name string) workspacesdk.AgentConn {
conn, err := tailnet.NewConn(&tailnet.Options{
Addresses: []netip.Prefix{tailnet.TailscaleServicePrefix.RandomPrefix()},
DERPMap: derpMap,
Logger: logger.Named(name),
})
require.NoError(t, err)
t.Cleanup(func() {
t.Logf("closing conn %s", name)
_ = conn.Close()
})
testCtx, testCtxCancel := context.WithCancel(context.Background())
t.Cleanup(testCtxCancel)
clientID := uuid.New()
ctrl := tailnet.NewTunnelSrcCoordController(logger, conn)
ctrl.AddDestination(agentID)
auth := tailnet.ClientCoordinateeAuth{AgentID: agentID}
coordination := ctrl.New(tailnet.NewInMemoryCoordinatorClient(logger, clientID, auth, coordinator))
t.Cleanup(func() {
t.Logf("closing coordination %s", name)
cctx, ccancel := context.WithTimeout(testCtx, testutil.WaitShort)
defer ccancel()
err := coordination.Close(cctx)
if err != nil {
t.Logf("error closing in-memory coordination: %s", err.Error())
}
t.Logf("closed coordination %s", name)
})
// Force DERP.
conn.SetBlockEndpoints(true)
sdkConn := workspacesdk.NewAgentConn(conn, workspacesdk.AgentConnOptions{
AgentID: agentID,
CloseFunc: func() error { return workspacesdk.ErrSkipClose },
})
t.Cleanup(func() {
t.Logf("closing sdkConn %s", name)
_ = sdkConn.Close()
})
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
if !sdkConn.AwaitReachable(ctx) {
t.Fatal("agent not reachable")
}
return sdkConn
}
conn1 := newClientConn(originalDerpMap, "client1")
// Change the DERP map.
newDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
require.NotNil(t, newDerpMap)
// Change the region ID.
newDerpMap.Regions[2] = newDerpMap.Regions[1]
delete(newDerpMap.Regions, 1)
newDerpMap.Regions[2].RegionID = 2
for _, node := range newDerpMap.Regions[2].Nodes {
node.RegionID = 2
}
// Push a new DERP map to the agent.
err := client.PushDERPMapUpdate(newDerpMap)
require.NoError(t, err)
t.Log("pushed DERPMap update to agent")
require.Eventually(t, func() bool {
conn := uut.TailnetConn()
if conn == nil {
return false
}
regionIDs := conn.DERPMap().RegionIDs()
preferredDERP := conn.Node().PreferredDERP
t.Logf("agent Conn DERPMap with regionIDs %v, PreferredDERP %d", regionIDs, preferredDERP)
return len(regionIDs) == 1 && regionIDs[0] == 2 && preferredDERP == 2
}, testutil.WaitLong, testutil.IntervalFast)
t.Log("agent got the new DERPMap")
// Connect from a second client and make sure it uses the new DERP map.
conn2 := newClientConn(newDerpMap, "client2")
require.Equal(t, []int{2}, conn2.TailnetConn().DERPMap().RegionIDs())
t.Log("conn2 got the new DERPMap")
// If the first client gets a DERP map update, it should be able to
// reconnect just fine.
conn1.TailnetConn().SetDERPMap(newDerpMap)
require.Equal(t, []int{2}, conn1.TailnetConn().DERPMap().RegionIDs())
t.Log("set the new DERPMap on conn1")
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
require.True(t, conn1.AwaitReachable(ctx))
t.Log("conn1 reached agent with new DERP")
}
func TestAgent_Speedtest(t *testing.T) {
t.Parallel()
t.Skip("This test is relatively flakey because of Tailscale's speedtest code...")
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
DERPMap: derpMap,
}, 0, func(client *agenttest.Client, options *agent.Options) {
options.Logger = logger.Named("agent")
})
defer conn.Close()
res, err := conn.Speedtest(ctx, speedtest.Upload, 250*time.Millisecond)
require.NoError(t, err)
t.Logf("%.2f MBits/s", res[len(res)-1].MBitsPerSecond())
}
func TestAgent_Reconnect(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
logger := testutil.Logger(t)
// After the agent is disconnected from a coordinator, it's supposed
// to reconnect!
fCoordinator := tailnettest.NewFakeCoordinator()
agentID := uuid.New()
statsCh := make(chan *proto.Stats, 50)
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
client := agenttest.NewClient(t,
logger,
agentID,
agentsdk.Manifest{
DERPMap: derpMap,
Directory: "/test/workspace",
},
statsCh,
fCoordinator,
)
defer client.Close()
closer := agent.New(agent.Options{
Client: client,
Logger: logger.Named("agent"),
})
defer closer.Close()
// Each iteration forces the agent to reconnect by closing
// the current coordinate call while the tracked HTTP server
// goroutine (from connection 1's createTailnet) is still
// alive, widening the race window.
const reconnections = 5
for i := range reconnections {
call := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
require.Equal(t, i+1, client.GetNumRefreshTokenCalls())
close(call.Resps) // hang up — triggers reconnect
}
// Verify final reconnect succeeds.
testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
require.Equal(t, reconnections+1, client.GetNumRefreshTokenCalls())
closer.Close()
}
func TestAgent_ReconnectNoLifecycleReemit(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
logger := testutil.Logger(t)
fCoordinator := tailnettest.NewFakeCoordinator()
agentID := uuid.New()
statsCh := make(chan *proto.Stats, 50)
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
client := agenttest.NewClient(t,
logger,
agentID,
agentsdk.Manifest{
DERPMap: derpMap,
Scripts: []codersdk.WorkspaceAgentScript{{
Script: "echo hello",
Timeout: 30 * time.Second,
RunOnStart: true,
}},
},
statsCh,
fCoordinator,
)
defer client.Close()
closer := agent.New(agent.Options{
Client: client,
Logger: logger.Named("agent"),
})
defer closer.Close()
// Wait for the agent to reach Ready state.
require.Eventually(t, func() bool {
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
}, testutil.WaitShort, testutil.IntervalFast)
statesBefore := slices.Clone(client.GetLifecycleStates())
// Disconnect by closing the coordinator response channel.
call1 := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
close(call1.Resps)
// Wait for reconnect.
testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
// Wait for a stats report as a deterministic steady-state proof.
testutil.RequireReceive(ctx, t, statsCh)
statesAfter := client.GetLifecycleStates()
require.Equal(t, statesBefore, statesAfter,
"lifecycle states should not be re-reported after reconnect")
closer.Close()
}
func TestAgent_WriteVSCodeConfigs(t *testing.T) {
t.Parallel()
logger := testutil.Logger(t)
coordinator := tailnet.NewCoordinator(logger)
defer coordinator.Close()
client := agenttest.NewClient(t,
logger,
uuid.New(),
agentsdk.Manifest{
GitAuthConfigs: 1,
DERPMap: &tailcfg.DERPMap{},
},
make(chan *proto.Stats, 50),
coordinator,
)
defer client.Close()
filesystem := afero.NewMemMapFs()
closer := agent.New(agent.Options{
Client: client,
Logger: logger.Named("agent"),
Filesystem: filesystem,
})
defer closer.Close()
home, err := os.UserHomeDir()
require.NoError(t, err)
name := filepath.Join(home, ".vscode-server", "data", "Machine", "settings.json")
require.Eventually(t, func() bool {
_, err := filesystem.Stat(name)
return err == nil
}, testutil.WaitShort, testutil.IntervalFast)
}
func TestAgent_DebugServer(t *testing.T) {
t.Parallel()
logDir := t.TempDir()
logPath := filepath.Join(logDir, "coder-agent.log")
randLogStr, err := cryptorand.String(32)
require.NoError(t, err)
require.NoError(t, os.WriteFile(logPath, []byte(randLogStr), 0o600))
newRotatedLogPath := filepath.Join(logDir, "coder-agent-2026-05-17T20-00-00.000.log")
oldRotatedLogPath := filepath.Join(logDir, "coder-agent-2026-05-17T19-00-00.000.log")
require.NoError(t, os.WriteFile(newRotatedLogPath, []byte("new rotated log"), 0o600))
require.NoError(t, os.WriteFile(oldRotatedLogPath, []byte("old rotated log"), 0o600))
now := time.Now()
newRotatedModTime := now.Add(-time.Minute)
oldRotatedModTime := now.Add(-48 * time.Hour)
require.NoError(t, os.Chtimes(newRotatedLogPath, newRotatedModTime, newRotatedModTime))
require.NoError(t, os.Chtimes(oldRotatedLogPath, oldRotatedModTime, oldRotatedModTime))
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
//nolint:dogsled
conn, _, _, _, agnt := setupAgentWithSecrets(t, agentsdk.Manifest{
DERPMap: derpMap,
EnvironmentVariables: map[string]string{
"AWS_SECRET_ACCESS_KEY": "env-value-should-be-redacted-67890",
"EMPTY_VAR": "",
},
}, []agentsdk.WorkspaceSecret{
{EnvName: "DEBUG_SECRET", Value: []byte("super-secret-value-12345")},
}, 0, func(c *agenttest.Client, o *agent.Options) {
o.LogDir = logDir
})
awaitReachableCtx := testutil.Context(t, testutil.WaitLong)
ok := conn.AwaitReachable(awaitReachableCtx)
require.True(t, ok)
_ = conn.Close()
srv := httptest.NewServer(agnt.HTTPDebug())
t.Cleanup(srv.Close)
t.Run("MagicsockDebug", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.Contains(t, string(resBody), "<h1>magicsock</h1>")
})
t.Run("MagicsockDebugLogging", func(t *testing.T) {
t.Parallel()
t.Run("Enable", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/t", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.Contains(t, string(resBody), "updated magicsock debug logging to true")
})
t.Run("Disable", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/0", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.Contains(t, string(resBody), "updated magicsock debug logging to false")
})
t.Run("Invalid", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/blah", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusBadRequest, res.StatusCode)
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.Contains(t, string(resBody), `invalid state "blah", must be a boolean`)
})
})
t.Run("Manifest", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
var v agentsdk.Manifest
require.NoError(t, json.NewDecoder(res.Body).Decode(&v))
require.NotNil(t, v)
})
t.Run("ManifestSecretsStripped", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
body, err := io.ReadAll(res.Body)
require.NoError(t, err)
// The response must not contain the secret value.
require.NotContains(t, string(body), "super-secret-value-12345")
// Confirm we can decode as a Manifest. The SDK type
// intentionally has no Secrets field, so there is nothing
// to leak through JSON encoding.
var v agentsdk.Manifest
require.NoError(t, json.Unmarshal(body, &v))
})
t.Run("ManifestEnvVarValuesRedacted", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusOK, res.StatusCode)
body, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.NotContains(t, string(body), "env-value-should-be-redacted-67890")
var v agentsdk.Manifest
require.NoError(t, json.Unmarshal(body, &v))
require.Contains(t, v.EnvironmentVariables, "AWS_SECRET_ACCESS_KEY")
require.Equal(t, "***REDACTED***", v.EnvironmentVariables["AWS_SECRET_ACCESS_KEY"])
// Empty values carry no secret and are preserved as empty.
require.Contains(t, v.EnvironmentVariables, "EMPTY_VAR")
require.Equal(t, "", v.EnvironmentVariables["EMPTY_VAR"])
})
t.Run("Logs", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/logs", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
require.Equal(t, http.StatusOK, res.StatusCode)
defer res.Body.Close()
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
require.NotEmpty(t, string(resBody))
require.Contains(t, string(resBody), randLogStr)
require.NotContains(t, string(resBody), "new rotated log")
})
t.Run("LogsIncludeActiveOnlyWithAfter", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
url := srv.URL + "/debug/logs?after=" + newRotatedModTime.Add(time.Minute).UTC().Format(time.RFC3339Nano)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
require.Equal(t, http.StatusOK, res.StatusCode)
defer res.Body.Close()
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
body := string(resBody)
require.Contains(t, body, randLogStr)
require.Contains(t, body, "coder-agent.log")
require.NotContains(t, body, "new rotated log")
require.NotContains(t, body, "old rotated log")
})
t.Run("LogsIncludeRotatedWithAfter", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
url := srv.URL + "/debug/logs?after=" + newRotatedModTime.Add(-time.Minute).UTC().Format(time.RFC3339Nano)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
require.Equal(t, http.StatusOK, res.StatusCode)
defer res.Body.Close()
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
body := string(resBody)
require.Contains(t, body, randLogStr)
require.Contains(t, body, "coder-agent.log")
require.Contains(t, body, "coder-agent-2026-05-17T20-00-00.000.log")
require.Contains(t, body, "new rotated log")
require.NotContains(t, body, "old rotated log")
})
t.Run("LogsIncludeRotatedWithOlderAfter", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
url := srv.URL + "/debug/logs?after=" + oldRotatedModTime.Add(-time.Minute).UTC().Format(time.RFC3339Nano)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
require.Equal(t, http.StatusOK, res.StatusCode)
defer res.Body.Close()
resBody, err := io.ReadAll(res.Body)
require.NoError(t, err)
body := string(resBody)
require.Contains(t, body, randLogStr)
require.Contains(t, body, "new rotated log")
require.Contains(t, body, "old rotated log")
require.Less(t, strings.Index(body, randLogStr), strings.Index(body, "new rotated log"))
require.Less(t, strings.Index(body, "new rotated log"), strings.Index(body, "old rotated log"))
})
t.Run("LogsInvalidAfter", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/logs?after=nope", nil)
require.NoError(t, err)
res, err := srv.Client().Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, http.StatusBadRequest, res.StatusCode)
})
}
func TestAgent_ScriptLogging(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("bash scripts only")
}
t.Parallel()
ctx := testutil.Context(t, testutil.WaitMedium)
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
logsCh := make(chan *proto.BatchCreateLogsRequest, 100)
lsStart := uuid.UUID{0x11}
lsStop := uuid.UUID{0x22}
//nolint:dogsled
_, _, _, _, agnt := setupAgent(
t,
agentsdk.Manifest{
DERPMap: derpMap,
Scripts: []codersdk.WorkspaceAgentScript{
{
LogSourceID: lsStart,
RunOnStart: true,
Script: `#!/bin/sh
i=0
while [ $i -ne 5 ]
do
i=$(($i+1))
echo "start $i"
done
`,
},
{
LogSourceID: lsStop,
RunOnStop: true,
Script: `#!/bin/sh
i=0
while [ $i -ne 3000 ]
do
i=$(($i+1))
echo "stop $i"
done
`, // send a lot of stop logs to make sure we don't truncate shutdown logs before closing the API conn
},
},
},
0,
func(cl *agenttest.Client, _ *agent.Options) {
cl.SetLogsChannel(logsCh)
},
)
n := 1
for n <= 5 {
logs := testutil.TryReceive(ctx, t, logsCh)
require.NotNil(t, logs)
for _, l := range logs.GetLogs() {
require.Equal(t, fmt.Sprintf("start %d", n), l.GetOutput())
n++
}
}
err := agnt.Close()
require.NoError(t, err)
n = 1
for n <= 3000 {
logs := testutil.TryReceive(ctx, t, logsCh)
require.NotNil(t, logs)
for _, l := range logs.GetLogs() {
require.Equal(t, fmt.Sprintf("stop %d", n), l.GetOutput())
n++
}
t.Logf("got %d stop logs", n-1)
}
}
// setupAgentSSHClient creates an agent, dials it, and sets up an ssh.Client for it
func setupAgentSSHClient(ctx context.Context, t *testing.T) *ssh.Client {
//nolint: dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := agentConn.SSHClient(ctx)
require.NoError(t, err)
t.Cleanup(func() { sshClient.Close() })
return sshClient
}
func setupSSHSession(
t *testing.T,
manifest agentsdk.Manifest,
banner codersdk.BannerConfig,
prepareFS func(fs afero.Fs),
opts ...func(*agenttest.Client, *agent.Options),
) *ssh.Session {
return setupSSHSessionOnPort(t, manifest, banner, prepareFS, workspacesdk.AgentSSHPort, opts...)
}
func setupSSHSessionOnPort(
t *testing.T,
manifest agentsdk.Manifest,
banner codersdk.BannerConfig,
prepareFS func(fs afero.Fs),
port uint16,
opts ...func(*agenttest.Client, *agent.Options),
) *ssh.Session {
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
opts = append(opts, func(c *agenttest.Client, o *agent.Options) {
c.SetAnnouncementBannersFunc(func() ([]codersdk.BannerConfig, error) {
return []codersdk.BannerConfig{banner}, nil
})
})
//nolint:dogsled
conn, _, _, fs, _ := setupAgent(t, manifest, 0, opts...)
if prepareFS != nil {
prepareFS(fs)
}
sshClient, err := conn.SSHClientOnPort(ctx, port)
require.NoError(t, err)
t.Cleanup(func() {
_ = sshClient.Close()
})
session, err := sshClient.NewSession()
require.NoError(t, err)
t.Cleanup(func() {
_ = session.Close()
})
return session
}
func setupAgent(t testing.TB, metadata agentsdk.Manifest, ptyTimeout time.Duration, opts ...func(*agenttest.Client, *agent.Options)) (
workspacesdk.AgentConn,
*agenttest.Client,
<-chan *proto.Stats,
afero.Fs,
agent.Agent,
) {
return setupAgentWithSecrets(t, metadata, nil, ptyTimeout, opts...)
}
// setupAgentWithSecrets is like setupAgent but also injects user
// secrets into the agent's proto manifest. Separate from setupAgent
// because agentsdk.Manifest intentionally does not carry secrets; see
// the Manifest doc comment in codersdk/agentsdk.
func setupAgentWithSecrets(t testing.TB, metadata agentsdk.Manifest, secrets []agentsdk.WorkspaceSecret, ptyTimeout time.Duration, opts ...func(*agenttest.Client, *agent.Options)) (
workspacesdk.AgentConn,
*agenttest.Client,
<-chan *proto.Stats,
afero.Fs,
agent.Agent,
) {
logger := slogtest.Make(t, &slogtest.Options{
// Agent can drop errors when shutting down, and some, like the
// fasthttplistener connection closed error, are unexported.
IgnoreErrors: true,
}).Leveled(slog.LevelDebug)
if metadata.DERPMap == nil {
metadata.DERPMap, _ = tailnettest.RunDERPAndSTUN(t)
}
if metadata.AgentID == uuid.Nil {
metadata.AgentID = uuid.New()
}
if metadata.AgentName == "" {
metadata.AgentName = "test-agent"
}
if metadata.WorkspaceName == "" {
metadata.WorkspaceName = "test-workspace"
}
if metadata.OwnerName == "" {
metadata.OwnerName = "test-user"
}
if metadata.WorkspaceID == uuid.Nil {
metadata.WorkspaceID = uuid.New()
}
coordinator := tailnet.NewCoordinator(logger)
t.Cleanup(func() {
_ = coordinator.Close()
})
statsCh := make(chan *proto.Stats, 50)
fs := afero.NewMemMapFs()
c := agenttest.NewClientWithSecrets(t, logger.Named("agenttest"), metadata.AgentID, metadata, secrets, statsCh, coordinator)
t.Cleanup(c.Close)
options := agent.Options{
Client: c,
Filesystem: fs,
Logger: logger.Named("agent"),
ReconnectingPTYTimeout: ptyTimeout,
EnvironmentVariables: map[string]string{},
StatsReportInterval: agenttest.StatsInterval,
}
for _, opt := range opts {
opt(c, &options)
}
agnt := agent.New(options)
t.Cleanup(func() {
_ = agnt.Close()
})
conn, err := tailnet.NewConn(&tailnet.Options{
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.TailscaleServicePrefix.RandomAddr(), 128)},
DERPMap: metadata.DERPMap,
Logger: logger.Named("client"),
})
require.NoError(t, err)
t.Cleanup(func() {
_ = conn.Close()
})
testCtx, testCtxCancel := context.WithCancel(context.Background())
t.Cleanup(testCtxCancel)
clientID := uuid.New()
ctrl := tailnet.NewTunnelSrcCoordController(logger, conn)
ctrl.AddDestination(metadata.AgentID)
auth := tailnet.ClientCoordinateeAuth{AgentID: metadata.AgentID}
coordination := ctrl.New(tailnet.NewInMemoryCoordinatorClient(
logger, clientID, auth, coordinator))
t.Cleanup(func() {
cctx, ccancel := context.WithTimeout(testCtx, testutil.WaitShort)
defer ccancel()
err := coordination.Close(cctx)
if err != nil {
t.Logf("error closing in-mem coordination: %s", err.Error())
}
})
agentConn := workspacesdk.NewAgentConn(conn, workspacesdk.AgentConnOptions{
AgentID: metadata.AgentID,
})
t.Cleanup(func() {
_ = agentConn.Close()
})
// Ideally we wouldn't wait too long here, but sometimes the the
// networking needs more time to resolve itself.
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
if !agentConn.AwaitReachable(ctx) {
t.Fatal("agent not reachable")
}
return agentConn, c, statsCh, fs, agnt
}
var dialTestPayload = []byte("dean-was-here123")
func testDial(ctx context.Context, t testing.TB, c net.Conn) {
t.Helper()
if deadline, ok := ctx.Deadline(); ok {
err := c.SetDeadline(deadline)
assert.NoError(t, err)
defer func() {
err := c.SetDeadline(time.Time{})
assert.NoError(t, err)
}()
}
assertWritePayload(t, c, dialTestPayload)
assertReadPayload(t, c, dialTestPayload)
}
func testAccept(ctx context.Context, t testing.TB, c net.Conn) {
t.Helper()
defer c.Close()
if deadline, ok := ctx.Deadline(); ok {
err := c.SetDeadline(deadline)
assert.NoError(t, err)
defer func() {
err := c.SetDeadline(time.Time{})
assert.NoError(t, err)
}()
}
assertReadPayload(t, c, dialTestPayload)
assertWritePayload(t, c, dialTestPayload)
}
func assertReadPayload(t testing.TB, r io.Reader, payload []byte) {
t.Helper()
b := make([]byte, len(payload)+16)
n, err := r.Read(b)
assert.NoError(t, err, "read payload")
assert.Equal(t, len(payload), n, "read payload length does not match")
assert.Equal(t, payload, b[:n])
}
func assertWritePayload(t testing.TB, w io.Writer, payload []byte) {
t.Helper()
n, err := w.Write(payload)
assert.NoError(t, err, "write payload")
assert.Equal(t, len(payload), n, "written payload length does not match")
}
func testSessionOutput(t *testing.T, session *ssh.Session, expected, unexpected []string, expectedRe *regexp.Regexp) {
t.Helper()
logger := testutil.Logger(t)
err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
var stdout bytes.Buffer
session.Stdout = &stdout
stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session)
err = session.Shell()
require.NoError(t, err)
stdin.WriteLine("exit 0")
waitErr := make(chan error, 1)
go func() {
waitErr <- session.Wait()
}()
select {
case err = <-waitErr:
require.NoError(t, err)
case <-time.After(testutil.WaitLong):
require.Fail(t, "timed out waiting for session to exit")
}
for _, unexpected := range unexpected {
require.NotContains(t, stdout.String(), unexpected, "should not show output")
}
for _, expect := range expected {
require.Contains(t, stdout.String(), expect, "should show output")
}
if expectedRe != nil {
require.Regexp(t, expectedRe, stdout.String())
}
}
func TestAgent_Metrics_SSH(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
registry := prometheus.NewRegistry()
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
o.PrometheusRegistry = registry
})
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
err = session.Shell()
require.NoError(t, err)
expected := []struct {
Name string
Type proto.Stats_Metric_Type
CheckFn func(float64) error
Labels []*proto.Stats_Metric_Label
}{
{
Name: "agent_reconnecting_pty_connections_total",
Type: proto.Stats_Metric_COUNTER,
CheckFn: func(v float64) error {
if v == 0 {
return nil
}
return xerrors.Errorf("expected 0, got %f", v)
},
},
{
Name: "agent_sessions_total",
Type: proto.Stats_Metric_COUNTER,
CheckFn: func(v float64) error {
if v == 1 {
return nil
}
return xerrors.Errorf("expected 1, got %f", v)
},
Labels: []*proto.Stats_Metric_Label{
{
Name: "magic_type",
Value: "ssh",
},
{
Name: "pty",
Value: "no",
},
},
},
{
Name: "agent_ssh_server_failed_connections_total",
Type: proto.Stats_Metric_COUNTER,
CheckFn: func(v float64) error {
if v == 0 {
return nil
}
return xerrors.Errorf("expected 0, got %f", v)
},
},
{
Name: "agent_ssh_server_sftp_connections_total",
Type: proto.Stats_Metric_COUNTER,
CheckFn: func(v float64) error {
if v == 0 {
return nil
}
return xerrors.Errorf("expected 0, got %f", v)
},
},
{
Name: "agent_ssh_server_sftp_server_errors_total",
Type: proto.Stats_Metric_COUNTER,
CheckFn: func(v float64) error {
if v == 0 {
return nil
}
return xerrors.Errorf("expected 0, got %f", v)
},
},
{
Name: "coderd_agentstats_currently_reachable_peers",
Type: proto.Stats_Metric_GAUGE,
CheckFn: func(float64) error {
// We can't reliably ping a peer here, and networking is out of
// scope of this test, so we just test that the metric exists
// with the correct labels.
return nil
},
Labels: []*proto.Stats_Metric_Label{
{
Name: "connection_type",
Value: "derp",
},
},
},
{
Name: "coderd_agentstats_currently_reachable_peers",
Type: proto.Stats_Metric_GAUGE,
CheckFn: func(float64) error {
return nil
},
Labels: []*proto.Stats_Metric_Label{
{
Name: "connection_type",
Value: "p2p",
},
},
},
{
Name: "coderd_agentstats_startup_script_seconds",
Type: proto.Stats_Metric_GAUGE,
CheckFn: func(f float64) error {
if f >= 0 {
return nil
}
return xerrors.Errorf("expected >= 0, got %f", f)
},
Labels: []*proto.Stats_Metric_Label{
{
Name: "success",
Value: "true",
},
},
},
}
var actual []*promgo.MetricFamily
assert.Eventually(t, func() bool {
actual, err = registry.Gather()
if err != nil {
return false
}
count := 0
for _, m := range actual {
count += len(m.GetMetric())
}
return count == len(expected)
}, testutil.WaitLong, testutil.IntervalFast)
i := 0
for _, mf := range actual {
for _, m := range mf.GetMetric() {
assert.Equal(t, expected[i].Name, mf.GetName())
assert.Equal(t, expected[i].Type.String(), mf.GetType().String())
switch expected[i].Type {
case proto.Stats_Metric_GAUGE:
assert.NoError(t, expected[i].CheckFn(m.GetGauge().GetValue()), "check fn for %s failed", expected[i].Name)
case proto.Stats_Metric_COUNTER:
assert.NoError(t, expected[i].CheckFn(m.GetCounter().GetValue()), "check fn for %s failed", expected[i].Name)
}
for j, lbl := range expected[i].Labels {
assert.Equal(t, m.GetLabel()[j], &promgo.LabelPair{
Name: &lbl.Name,
Value: &lbl.Value,
})
}
i++
}
}
_, err = stdin.Write([]byte("exit 0\n"))
require.NoError(t, err, "writing exit to stdin")
_ = stdin.Close()
err = session.Wait()
require.NoError(t, err, "waiting for session to exit")
}
// echoOnce accepts a single connection, reads 4 bytes and echos them back
func echoOnce(t *testing.T, ll net.Listener) {
t.Helper()
conn, err := ll.Accept()
if err != nil {
return
}
defer conn.Close()
b := make([]byte, 4)
_, err = conn.Read(b)
if !assert.NoError(t, err) {
return
}
_, err = conn.Write(b)
if !assert.NoError(t, err) {
return
}
}
// requireEcho sends 4 bytes and requires the read response to match what was
// sent. Returns the number of bytes written.
func requireEcho(t *testing.T, conn net.Conn) int {
t.Helper()
written, err := conn.Write([]byte("test"))
require.NoError(t, err)
b := make([]byte, written)
_, err = conn.Read(b)
require.NoError(t, err)
require.Equal(t, "test", string(b))
return written
}
const (
wantSomeBytes int64 = -1
wantIgnoreBytes int64 = -2
// wantIgnoreCode skips the disconnect status code check, for disconnects
// whose code depends on process timing.
wantIgnoreCode = -1
)
// assertConnectionReport expects the provided connect and disconnect events
// have been reported. If the disconnection bytes are wantSomeBytes, assert the
// bytes are > 0, if wantIgnoreBytes do not check the bytes, otherwise assert
// they match exactly. If the disconnect code is wantIgnoreCode, do not check
// the status code, otherwise assert it matches exactly.
func assertConnectionReport(t testing.TB, agentClient *agenttest.Client,
connect proto.ConnectEvent, disconnect proto.DisconnectEvent,
) {
t.Helper()
var reports []*proto.ReportConnectionRequest
if !assert.Eventually(t, func() bool {
reports = agentClient.GetConnectionReports()
return len(reports) >= 2
}, testutil.WaitMedium, testutil.IntervalFast, "waiting for 2 connection reports or more; got %d", len(reports)) {
return
}
assert.Len(t, reports, 2, "want 2 connection reports")
assert.Equal(t, proto.Connection_CONNECT, reports[0].GetConnection().GetAction(), "first report should be connect")
assert.Equal(t, proto.Connection_DISCONNECT, reports[1].GetConnection().GetAction(), "second report should be disconnect")
assert.Equal(t, connect.Method, reports[0].GetConnection().GetConnectionMethod(), "connect method should be %s", connect.Method)
assert.Equal(t, connect.Method, reports[1].GetConnection().GetConnectionMethod(), "disconnect method should be %s", connect.Method)
assert.Equal(t, connect.AppName, reports[0].GetConnection().GetAppName(), "connect app name should be %s", connect.AppName)
assert.Equal(t, connect.AppName, reports[1].GetConnection().GetAppName(), "disconnect app name should be %s", connect.AppName)
assert.Equal(t, reports[0].GetConnection().GetId(), reports[1].GetConnection().GetId(), "disconnect id should match connect id")
t1 := reports[0].GetConnection().GetTimestamp().AsTime()
t2 := reports[1].GetConnection().GetTimestamp().AsTime()
assert.True(t, t1.Before(t2) || t1.Equal(t2), "connect timestamp should be before or equal to disconnect timestamp")
assert.NotEmpty(t, reports[0].GetConnection().GetIp(), "connect ip should not be empty")
assert.NotEmpty(t, reports[1].GetConnection().GetIp(), "disconnect ip should not be empty")
assert.Equal(t, 0, int(reports[0].GetConnection().GetStatusCode()), "connect status code should be 0")
if disconnect.Code != wantIgnoreCode {
assert.Equal(t, disconnect.Code, int(reports[1].GetConnection().GetStatusCode()), "disconnect status code should be %d", disconnect.Code)
}
assert.Equal(t, "", reports[0].GetConnection().GetReason(), "connect reason should be empty")
if disconnect.Reason != "" {
assert.Contains(t, reports[1].GetConnection().GetReason(), disconnect.Reason, "disconnect reason should contain %s", disconnect.Reason)
} else {
assert.Equal(t, "", reports[1].GetConnection().GetReason(), "disconnect reason should be empty")
}
assert.Equal(t, connect.ClientSessionID, reports[0].GetConnection().GetClientSessionId(), "connect client session id should be %s", connect.ClientSessionID)
assert.Equal(t, connect.ClientSessionID, reports[1].GetConnection().GetClientSessionId(), "disconnect client session id should be %s", connect.ClientSessionID)
assert.Equal(t, int64(0), reports[0].GetConnection().GetRxBytes(), "connect rxbytes should be 0")
assert.Equal(t, int64(0), reports[0].GetConnection().GetTxBytes(), "connect txbytes should be 0")
assertBytes(t, disconnect.RxBytes, reports[1].GetConnection().GetRxBytes(), "rx")
assertBytes(t, disconnect.TxBytes, reports[1].GetConnection().GetTxBytes(), "tx")
}
func assertBytes(t testing.TB, expectedBytes, gotBytes int64, label string) {
t.Helper()
if expectedBytes == wantSomeBytes {
assert.Greater(t, gotBytes, int64(0), "disconnect %s bytes should be greater than zero", label)
} else if expectedBytes != wantIgnoreBytes {
assert.Equal(t, expectedBytes, gotBytes, "disconnect %s bytes should be %d", label, expectedBytes)
}
}
func TestAgent_ToolCall(t *testing.T) {
t.Skip("flaky until PLAT-717: the agent netstack drops a SYN that reuses a source port it holds in TIME_WAIT")
t.Parallel()
//nolint:dogsled
conn, _, _, agentFS, _ := setupAgent(t, agentsdk.Manifest{}, 0)
conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
// Every call uses the tool call context, as chatd does, so a method
// that wrongly sent the tool call ID would get a saved response.
tests := []struct {
name string
fs afero.Fs
act func(ctx context.Context, path string) (any, error)
saved func(workspacesdk.CancelToolCallResponse) (any, error)
}{
{
name: "StartProcess",
fs: afero.NewOsFs(),
act: func(ctx context.Context, path string) (any, error) {
resp, err := conn.StartProcess(ctx, workspacesdk.StartProcessRequest{
Command: "printf run > " + filepath.Base(path),
WorkDir: filepath.Dir(path),
})
if err != nil {
return nil, err
}
out, err := conn.ProcessOutput(ctx, resp.ID, &workspacesdk.ProcessOutputOptions{Wait: true})
if err != nil {
return nil, err
}
if out.ExitCode == nil {
return nil, xerrors.Errorf("process output reports no exit: %+v", out)
}
return resp, nil
},
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.StartProcessResult() },
},
{
name: "WriteFile",
fs: agentFS,
act: func(ctx context.Context, path string) (any, error) {
return nil, conn.WriteFile(ctx, path, strings.NewReader("run"))
},
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return nil, r.WriteFileResult() },
},
{
name: "EditFiles",
fs: agentFS,
act: func(ctx context.Context, path string) (any, error) {
return conn.EditFiles(ctx, workspacesdk.FileEditRequest{
Files: []workspacesdk.FileEdits{{Path: path, Edits: []workspacesdk.FileEdit{{OldText: "reset", NewText: "run"}}}},
IncludeDiff: true,
})
},
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.EditFilesResult() },
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
path := filepath.Join(t.TempDir(), "tool-call.txt")
reset := func() { require.NoError(t, afero.WriteFile(tc.fs, path, []byte("reset"), 0o600)) }
requireContent := func(want string) {
t.Helper()
content, err := afero.ReadFile(tc.fs, path)
require.NoError(t, err)
require.Equal(t, want, string(content))
}
id := uuid.New()
toolCtx := workspacesdk.WithToolCallID(ctx, id)
reset()
first, err := tc.act(toolCtx, path)
require.NoError(t, err)
requireContent("run")
reset()
repeat, err := tc.act(toolCtx, path)
require.NoError(t, err)
require.Equal(t, first, repeat)
requireContent("reset")
canceled, err := conn.CancelToolCall(toolCtx, id)
require.NoError(t, err)
require.True(t, canceled.Received)
saved, err := tc.saved(canceled)
require.NoError(t, err)
require.Equal(t, first, saved)
_, err = tc.act(toolCtx, path)
var sdkErr *codersdk.Error
require.ErrorAs(t, err, &sdkErr)
require.Equal(t, http.StatusConflict, sdkErr.StatusCode())
})
}
}
// TestAgent_ToolCallOtherChat checks that two chats using the same tool
// call ID get separate processes and cancels.
func TestAgent_ToolCallOtherChat(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
agentID := uuid.New()
//nolint:dogsled
connA, _, _, _, _ := setupAgent(t, agentsdk.Manifest{AgentID: agentID}, 0)
connA.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
// connB shares connA's tailnet connection, which setupAgent closes.
connB := workspacesdk.NewAgentConn(connA.TailnetConn(), workspacesdk.AgentConnOptions{AgentID: agentID})
connB.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
id := uuid.New()
toolCtx := workspacesdk.WithToolCallID(ctx, id)
_, err := connA.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300"})
require.NoError(t, err)
t.Cleanup(func() {
_, err := connA.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id)
assert.NoError(t, err)
})
_, err = connB.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "printf b"})
require.NoError(t, err)
outB, err := connB.ProcessOutput(toolCtx, id.String(), &workspacesdk.ProcessOutputOptions{Wait: true})
require.NoError(t, err)
require.Equal(t, "printf b", outB.Command)
require.Equal(t, "b", outB.Output)
_, err = connB.CancelToolCall(toolCtx, id)
require.NoError(t, err)
// CancelToolCall marks the process canceled before returning, so a
// wrong cancel shows here without waiting.
outA, err := connA.ProcessOutput(toolCtx, id.String(), nil)
require.NoError(t, err)
require.Equal(t, "sleep 300", outA.Command)
require.False(t, outA.Canceled)
}
func TestAgent_ProcessOutputTimeoutFromStartProcess(t *testing.T) {
t.Parallel()
//nolint:dogsled
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
// A 1 ms timeout has passed before any wait uses it and one hour cannot
// pass during the test, so no case depends on timing.
tests := []struct {
name string
timeoutMs int64
opts workspacesdk.ProcessOutputOptions
cancel bool
want workspacesdk.ProcessOutputResponse
}{
{
name: "TimeoutFromStartProcess",
timeoutMs: 1,
opts: workspacesdk.ProcessOutputOptions{Wait: true, TimeoutFromStartProcess: true},
want: workspacesdk.ProcessOutputResponse{Running: true, TimedOut: true},
},
{
name: "PlainWait",
timeoutMs: 1,
opts: workspacesdk.ProcessOutputOptions{Wait: true},
cancel: true,
want: workspacesdk.ProcessOutputResponse{Canceled: true},
},
{
name: "BeforeTimeout",
timeoutMs: time.Hour.Milliseconds(),
want: workspacesdk.ProcessOutputResponse{Running: true},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
id := uuid.New()
toolCtx := workspacesdk.WithToolCallID(ctx, id)
_, err := conn.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300", TimeoutMs: tc.timeoutMs})
require.NoError(t, err)
// Cancel the tool call so the agent does not wait for its process
// on close.
t.Cleanup(func() {
_, err := conn.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id)
assert.NoError(t, err)
})
type result struct {
resp workspacesdk.ProcessOutputResponse
err error
}
waited := make(chan result, 1)
go func() {
resp, err := conn.ProcessOutput(toolCtx, id.String(), &tc.opts)
waited <- result{resp, err}
}()
if tc.cancel {
_, err := conn.CancelToolCall(toolCtx, id)
require.NoError(t, err)
}
got := testutil.RequireReceive(ctx, t, waited)
require.NoError(t, got.err)
require.Equal(t, tc.want.Running, got.resp.Running)
require.Equal(t, tc.want.TimedOut, got.resp.TimedOut)
require.Equal(t, tc.want.Canceled, got.resp.Canceled)
})
}
}